diff --git a/CLAUDE.md b/CLAUDE.md
index b0e7c095..453392f2 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -131,6 +131,69 @@ The hierarchical Bayesian inference engine. Key files:
See **YAML Standards** below for the full conventions. The tfmodel config (`tfs_configure_config.yaml`) is generated by `tfs-configure-model` and drives `tfs-fit-model`. It contains `data`, `components`, `priors_file`, and `guesses_file` sections. Do not hand-edit it.
+## Simulation Internals (`simulate/`)
+
+### Overview
+
+The simulation pipeline generates ground-truth phenotypes for a synthetic TF screen experiment. The top-level entry point is `library_prediction` (`simulate/library_prediction.py`), which returns five dataframes:
+
+| Return value | Content |
+|---|---|
+| `library_df` | One row per genotype in the library |
+| `phenotype_df` | Long-form growth predictions (one row per genotype × condition) |
+| `genotype_theta_df` | Long-form theta predictions (one row per genotype × titrant_conc) |
+| `parameters_df` | One row per genotype; per-genotype Hill/theta params + dk_geno + activity |
+| `binding_theta_df` | Theta at binding concentrations for calibration genotypes (`None` if not configured) |
+
+The core calculation happens in `thermo_to_growth` (`simulate/thermo_to_growth.py`):
+
+1. **Prior-predictive theta sampling** — `sample_theta_prior` draws a `theta_gc` matrix of shape `(G, C)` where G = number of library genotypes and C = number of unique titrant concentrations. The genotype order here matches the library (`sim_data`) order.
+2. **theta_gc_override injection** — specific rows of `theta_gc` can be replaced before any further computation (see below).
+3. **Growth rate calculation** — theta is mapped to growth rates via the configured `growth_params`.
+4. **parameters_df assembly** — per-genotype Hill parameters extracted from the `theta_param` pytree, then patched with `theta_params_override` (see below).
+
+### Binding data and calibration genotypes
+
+The optional `binding_data` YAML block configures calibration genotypes for which measured binding curves are available. It has two sub-paths that can coexist:
+
+**Path 1: Simulated genotypes** (`binding_data.genotypes` list)
+- `wt` gets its natural unperturbed prior-predictive reference curve.
+- Non-wt entries are drawn via a stratified (greedy maximin) algorithm across binding concentrations, ensuring diverse coverage.
+- The selected theta values at *growth* concentrations are injected via `theta_gc_override` so that the simulated growth data is consistent with the binding data.
+
+**Path 2: Measured Hill parameters** (`binding_data.genotype_params_file`)
+- Reads a CSV with columns `genotype, theta_low, theta_high, log_hill_K, hill_n`.
+- Only supported for Hill-based theta components (`hill_geno`, `hill_mut`).
+- **`theta_low` and `theta_high` are clamped to `[1e-4, 1-1e-4]` at read time** with a `UserWarning`. This is critical: a value of e.g. `1.000004` (a common float-rounding artefact) maps to `logit ≈ +16`, making all per-mutation deltas ~13–15 σ under the `HalfNormal(1)` prior on delta scales and preventing inference from recovering reasonable theta values.
+- For `hill_mut`: the function `build_theta_gc_override_hill_mut` assembles theta for **all library genotypes** (not just the measured ones) by additively combining per-mutation logit-space deltas. The WT reference is taken from `SimPriors` defaults. Multi-mutant genotypes not directly measured in the CSV are assembled from single-mutant deltas; directly-measured multi-mutants use their CSV values directly.
+- Measured genotypes override any earlier simulated-path values in `theta_gc_override`.
+
+### theta_gc_override and theta_params_override
+
+These two dicts are the mechanism by which binding data is "pinned" into the growth simulation.
+
+**`theta_gc_override`** (`dict[str, np.ndarray]`):
+- Keys are genotype strings; values are 1-D arrays of theta at the growth titrant concentrations (sorted-unique order from `sample_df`).
+- Applied in `thermo_to_growth` *after* prior-predictive sampling but *before* noise and all downstream computations.
+- Overwrites the corresponding row of `theta_gc` in-place using `geno_to_sim_idx` (which maps genotype → its index in the original library list).
+
+**`theta_params_override`** (`dict[str, dict[str, float]]`):
+- Keys are genotype strings; values are dicts with Hill parameter keys (`theta_low`, `theta_high`, `log_hill_K`, `hill_n`).
+- Applied in `thermo_to_growth` *after* `parameters_df` is assembled from the `theta_param` pytree, overwriting the relevant columns.
+- Purpose: the `theta_param` from `sample_theta_prior` reflects the prior-predictive draw, not the override. Without this patch, `tfs_sim_parameters.csv` would show the pre-override values, making the saved parameters inconsistent with what was actually simulated.
+- Only columns already present in `parameters_df` are updated; unknown keys are silently skipped.
+- Genotype keys not found in `parameters_df` are silently skipped.
+
+### Key files
+
+| File | Role |
+|---|---|
+| `simulate/library_prediction.py` | Top-level orchestrator; assembles override dicts and calls `thermo_to_growth` |
+| `simulate/thermo_to_growth.py` | Prior-predictive sampling, theta injection, growth calculation, parameters_df assembly |
+| `simulate/binding_params.py` | CSV reading (with theta clipping), per-component override builders, binding theta assemblers |
+| `simulate/sample_theta.py` | `sample_theta_prior` (prior-predictive) and `sample_theta_stratified` (greedy maximin) |
+| `simulate/sim_data_class.py` | `SimData` container and `build_sim_data` factory |
+
## YAML Standards
All YAML files in this codebase follow these conventions. Apply them when creating or modifying any YAML file.
diff --git a/docs/badges/coverage-badge.svg b/docs/badges/coverage-badge.svg
index b213db1a..92cb9a75 100644
--- a/docs/badges/coverage-badge.svg
+++ b/docs/badges/coverage-badge.svg
@@ -1 +1 @@
-
\ No newline at end of file
+
\ No newline at end of file
diff --git a/docs/badges/tests-badge.svg b/docs/badges/tests-badge.svg
index 14389206..c9b16aaa 100644
--- a/docs/badges/tests-badge.svg
+++ b/docs/badges/tests-badge.svg
@@ -1 +1 @@
-
\ No newline at end of file
+
\ No newline at end of file
diff --git a/docs/source/_static/saa_growth_calibration.png b/docs/source/_static/saa_growth_calibration.png
new file mode 100644
index 00000000..13b2c1c7
Binary files /dev/null and b/docs/source/_static/saa_growth_calibration.png differ
diff --git a/docs/source/_static/saa_growth_corr.png b/docs/source/_static/saa_growth_corr.png
new file mode 100644
index 00000000..40fe9373
Binary files /dev/null and b/docs/source/_static/saa_growth_corr.png differ
diff --git a/docs/source/_static/saa_params_log_hill_K.png b/docs/source/_static/saa_params_log_hill_K.png
new file mode 100644
index 00000000..6ca560c4
Binary files /dev/null and b/docs/source/_static/saa_params_log_hill_K.png differ
diff --git a/docs/source/_static/saa_theta_corr.png b/docs/source/_static/saa_theta_corr.png
new file mode 100644
index 00000000..6e058e9b
Binary files /dev/null and b/docs/source/_static/saa_theta_corr.png differ
diff --git a/docs/source/_static/saa_wt_theta_fits.png b/docs/source/_static/saa_wt_theta_fits.png
new file mode 100644
index 00000000..5294bd9c
Binary files /dev/null and b/docs/source/_static/saa_wt_theta_fits.png differ
diff --git a/docs/source/analysis.rst b/docs/source/analysis.rst
index 76e8582e..b9c25da4 100644
--- a/docs/source/analysis.rst
+++ b/docs/source/analysis.rst
@@ -53,8 +53,15 @@ conditions present in both the growth and binding data.
After convergence the script updates the production ``{out_prefix}_priors.csv``
and ``{out_prefix}_guesses.csv`` files in place (a ``.bak`` backup is written
-first), giving the full production fit a warm start. Diagnostic PDFs (one per
-genotype) and a ``{out_prefix}_calib_stats.json`` file are also written.
+first), giving the full production fit a warm start.
+
+Diagnostic artefacts from the calibration MAP run are also written (default
+prefix ``tfs_prefit``):
+
+* ``tfs_prefit_params.npz`` — MAP point estimates in constrained space.
+* ``tfs_prefit_checkpoint.pkl`` — optimizer checkpoint (can be passed to
+ ``--checkpoint_file`` to resume an interrupted calibration run).
+* ``tfs_prefit_losses.txt`` — per-epoch loss history.
.. code-block:: bash
@@ -139,9 +146,32 @@ writes one CSV per parameter group.
Outputs (default ``--out_prefix tfs_params``):
-* ``tfs_params_activity.csv`` — per-genotype TF activity *A* with posterior quantiles
-* ``tfs_params_theta.csv`` — inferred occupancy parameters (Hill *Kd*, *n*, etc.)
-* Additional CSVs depending on the components selected during configuration.
+* ``tfs_params_log_hill_K.csv`` — per-genotype log₁₀(*Kd*) with posterior quantiles.
+* ``tfs_params_theta_low.csv``, ``tfs_params_theta_high.csv`` — per-genotype
+ lower and upper occupancy plateaux.
+* ``tfs_params_dk_geno.csv`` — per-genotype pleiotropic growth effect.
+* Additional CSVs depending on the components selected during configuration
+ (e.g. ``tfs_params_d_log_hill_K.csv`` for per-mutation effects when using
+ ``hill_mut``; ``tfs_params_epi_*.csv`` for epistasis terms when
+ ``--epistasis`` is set).
+
+**Quantile columns**
+
+Each CSV contains metadata columns (typically ``genotype`` and
+``titrant_name``) followed by 17 quantile columns named ``q{level}``:
+
+.. code-block:: text
+
+ q0.001 q0.005 q0.01 q0.025 q0.05 q0.1 q0.159 q0.25
+ q0.5 q0.75 q0.841 q0.9 q0.95 q0.975 q0.99 q0.995 q0.999
+
+``q0.5`` is the posterior median (the recommended point estimate). ``q0.025``
+and ``q0.975`` span the 95% credible interval; ``q0.159`` and ``q0.841``
+span the ±1 σ interval under a normal approximation.
+
+When ``tfs-extract-params`` is given a MAP checkpoint (``.pkl``) instead of
+a posterior file, only a single ``q0.5`` column is written (the MAP point
+estimate, not a true posterior quantile).
Step 6: Predict Growth (``tfs-predict-growth``)
------------------------------------------------
@@ -206,6 +236,22 @@ Output (default ``--out_prefix tfs_cat_response``):
* ``tfs_cat_response.csv`` — one row per (genotype, titrant_name) with
``best_model``, AIC weights, and fitted parameters for every model.
+Step 9: Summarise Fit (``tfs-summarize-fit``)
+----------------------------------------------
+
+Collects the outputs of the completed run, computes prediction quality
+statistics, and writes diagnostic plots and tables to a ``summary/``
+subdirectory. It can be run after Steps 5–7 are complete and does not
+require the full posterior file to be retained.
+
+.. code-block:: bash
+
+ tfs-summarize-fit out/
+
+For a full guide to every output file — including how to read the theta
+correlation plots, calibration curves, and parameter-recovery scatter
+plots — see :doc:`summarize-fit`.
+
---
.. _model-components:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 394e858c..ad036cab 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -8,9 +8,11 @@ tfscreen
:maxdepth: 2
:caption: Contents:
+ quickstart
simulation
process-raw
analysis
+ summarize-fit
grid
ligandmpnn-features
@@ -30,6 +32,12 @@ The thermodynamic observable is operator occupancy (*θ*), inferred from *E. col
The model in ``tfscreen.tfmodel`` is designed to jointly capture high-throughput growth data and low-throughput binding data, learning the quantitative relationship between operator occupancy and growth rate in the process.
+Quick Start
+-----------
+New to ``tfscreen``? The :doc:`quickstart` page walks through a complete
+simulate-and-analyze run using the bundled example in
+``examples/simulate-and-analyze/``.
+
Simulation
----------
`tfscreen` allows you to simulate high-throughput screens starting from thermodynamic models of transcription factor binding and activity. See the :doc:`simulation` page for more details.
diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
new file mode 100644
index 00000000..6d8483fc
--- /dev/null
+++ b/docs/source/quickstart.rst
@@ -0,0 +1,156 @@
+===========
+Quick Start
+===========
+
+This page walks through a complete simulate-and-analyze run using the
+bundled example in ``examples/simulate-and-analyze/``. By the end you
+will have simulated a synthetic high-throughput TF-library screen, fitted
+the hierarchical Bayesian model, and produced a full set of diagnostic
+plots and statistics.
+
+The example uses a small library (~600 genotypes across three sub-libraries)
+so the full pipeline finishes in roughly 20–60 minutes on a laptop CPU or a
+few minutes on a GPU.
+
+Prerequisites
+-------------
+
+Install ``tfscreen`` and verify the entry points are on your PATH:
+
+.. code-block:: bash
+
+ git clone https://github.com/harmslab/tfscreen
+ cd tfscreen
+ pip install -e .
+ tfs-simulate --help # should print usage
+
+The example requires JAX and Numpyro, which are installed as dependencies.
+On Apple Silicon you may need to install the Metal-accelerated ``jax-metal``
+package separately; on a Linux cluster with a GPU, install the appropriate
+``jaxlib`` wheel for your CUDA version.
+
+Getting the example
+-------------------
+
+Copy the example directory to a working location:
+
+.. code-block:: bash
+
+ cp -r examples/simulate-and-analyze/ ~/tfscreen-example
+ cd ~/tfscreen-example
+
+The directory contains:
+
+* ``simulate_config.yaml`` — simulation parameters (library genetics,
+ growth conditions, binding data).
+* ``hill_params.csv`` — per-genotype Hill binding parameters that set the
+ ground-truth θ values used during simulation.
+* ``run.sh`` — the pipeline script.
+
+Running the pipeline
+--------------------
+
+.. code-block:: bash
+
+ bash run.sh simulate_config.yaml out/ 1
+
+The three positional arguments are:
+
+1. The simulation config file.
+2. The output directory (created automatically).
+3. The random seed (``1`` reproduces the documented example outputs).
+
+On a local machine the script uses JAX's CPU parallelism
+(``XLA_FLAGS="--xla_force_host_platform_device_count=8"``). On a
+cluster, comment that line out and uncomment ``module load cuda/...``
+instead.
+
+The script runs nine steps in sequence, printing a ``>>>`` header before
+each one:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 5 30 65
+
+ * - Step
+ - Command
+ - What it does
+ * - 1
+ - ``tfs-simulate``
+ - Builds the genotype library, draws per-genotype θ from the Hill
+ model, converts θ to growth rates, and writes simulated sequencing
+ data.
+ * - 2
+ - ``tfs-configure-model``
+ - Validates inputs, selects model components, writes
+ ``tfs_configure_config.yaml``.
+ * - 3
+ - ``tfs-prefit-calibration``
+ - Fast MAP fit to calibrate the growth-linking-function priors.
+ * - 4
+ - ``tfs-fit-model``
+ - Main SVI inference; writes ``tfs_fit_model_checkpoint.pkl``.
+ * - 5
+ - ``tfs-sample-posterior``
+ - Draws posterior samples; writes ``tfs_posterior.h5`` (~5 GB).
+ * - 6
+ - ``tfs-extract-params``
+ - Summarises the posterior into per-parameter CSV files.
+ * - 7
+ - ``tfs-predict-theta``
+ - Predicts θ at every (genotype, titrant concentration) in the
+ training data.
+ * - 8
+ - ``tfs-predict-growth``
+ - Predicts ln(CFU) with posterior quantiles for all training
+ observations.
+ * - 9
+ - ``tfs-summarize-fit``
+ - Computes fit statistics and writes diagnostic plots and CSVs to
+ ``out/summary/``.
+
+Expected outputs
+----------------
+
+After the run completes, ``out/`` will contain:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 38 62
+
+ * - File / directory
+ - Contents
+ * - ``tfs_sim_library.csv``
+ - All genotypes in the simulated library.
+ * - ``tfs_sim_growth.csv``
+ - Simulated ln(CFU) data (same format as ``tfs-process-counts`` output).
+ * - ``tfs_sim_binding.csv``
+ - Simulated binding curve observations.
+ * - ``tfs_sim_parameters.csv``
+ - Ground-truth per-genotype parameters (θ, dk_geno, etc.).
+ * - ``tfs_configure_config.yaml``
+ - Model configuration read by all downstream steps.
+ * - ``tfs_fit_model_checkpoint.pkl``
+ - Fitted model checkpoint.
+ * - ``tfs_posterior.h5``
+ - Posterior samples (~5 GB; not committed to the repo).
+ * - ``tfs_params_*.csv``
+ - Posterior summaries (quantiles) for each parameter group.
+ * - ``tfs_pred_theta.csv``
+ - Predicted θ with posterior quantiles.
+ * - ``tfs_pred_growth.csv``
+ - Predicted ln(CFU) with posterior quantiles.
+ * - ``summary/``
+ - Diagnostic plots and statistics from ``tfs-summarize-fit``.
+
+Next steps
+----------
+
+* **Understand each analysis step** — see :doc:`analysis` for the full CLI
+ reference covering every ``tfs-*`` command in the pipeline.
+* **Interpret the diagnostic outputs** — see :doc:`summarize-fit` for a
+ guided walkthrough of every plot and statistic in ``out/summary/``.
+* **Simulate a parameter sweep** — see :doc:`grid` for how to run a
+ Cartesian grid of configurations.
+* **Use real data** — see :doc:`process-raw` for converting FASTQ reads
+ into the ``growth.csv`` format that ``tfs-configure-model`` accepts.
diff --git a/docs/source/summarize-fit.rst b/docs/source/summarize-fit.rst
new file mode 100644
index 00000000..25315f70
--- /dev/null
+++ b/docs/source/summarize-fit.rst
@@ -0,0 +1,334 @@
+=================
+tfs-summarize-fit
+=================
+
+``tfs-summarize-fit`` collects the outputs of a completed model-fitting run,
+computes prediction quality statistics, and writes a set of diagnostic plots
+and tables to a ``summary/`` subdirectory.
+
+.. note::
+
+ **CSV files are the authoritative outputs.** Every PDF written by this
+ command has a matching CSV containing the exact data used to generate it.
+ Use the CSVs for downstream quantitative analysis; treat the PDFs as
+ human-readable summaries.
+
+Usage
+-----
+
+.. code-block:: bash
+
+ tfs-summarize-fit [options]
+
+**Positional arguments:**
+
+* ``run_dir`` — directory containing the completed model-fit outputs (i.e.
+ the ``out/`` directory produced by ``run.sh``).
+
+**Optional arguments:**
+
+* ``--ref_theta_file`` — CSV with columns ``genotype``, ``titrant_name``,
+ ``titrant_conc``, and a ``theta_obs`` or ``theta`` column containing
+ known θ values for evaluating out-of-sample predictions. When omitted,
+ the script automatically looks for ``*_sim_genotype_theta.csv`` in
+ ``run_dir`` (the file written by ``tfs-simulate``).
+* ``--out_prefix`` — prefix for all output files. Defaults to
+ ``{run_dir}/summary/tfs_summarize``.
+
+What the script needs
+---------------------
+
+The script scans ``run_dir`` for:
+
+* ``*_config.yaml`` — model configuration (required).
+* ``*_pred_theta.csv`` — theta predictions from ``tfs-predict-theta`` (required).
+* ``*_sim_genotype_theta.csv`` — ground-truth θ from ``tfs-simulate``
+ (used for test statistics; absent for real-data runs).
+* ``*_pred_growth.csv`` — growth predictions from ``tfs-predict-growth``
+ (optional; enables growth statistics and plots).
+* ``*_losses.txt`` — training loss history from ``tfs-fit-model``
+ (optional; enables the loss-curve plot).
+* ``*_params_*.csv`` — parameter summaries from ``tfs-extract-params``
+ (optional; enables parameter-recovery plots when ``*_sim_parameters.csv``
+ is also present).
+
+Output file reference
+---------------------
+
+.. list-table::
+ :header-rows: 1
+ :widths: 45 55
+
+ * - File
+ - Contents
+ * - ``tfs_summarize_fit_summary.json``
+ - Nested statistics dict (see :ref:`fit-summary-json`).
+ * - ``tfs_summarize_theta_corr.pdf`` / ``.csv``
+ - Two-panel θ correlation: training (left) and test (right).
+ * - ``tfs_summarize_theta_corr_training.csv``
+ - Joined (ref, predicted) θ pairs for training genotypes.
+ * - ``tfs_summarize_theta_corr_test.csv``
+ - Same for test genotypes.
+ * - ``tfs_summarize_growth_corr.pdf`` / ``.csv``
+ - Observed vs predicted ln(CFU) across all training points.
+ * - ``tfs_summarize_{genotype}_theta_fits.pdf`` / ``.csv``
+ - Per-genotype θ curve overlaid with binding observations.
+ * - ``tfs_summarize_{genotype}_trajectory.pdf`` / ``.csv``
+ - Per-genotype predicted ln(CFU) vs time across all conditions.
+ * - ``tfs_summarize_losses.pdf``
+ - Training loss curve.
+ * - ``tfs_summarize_theta_training_calibration.pdf``
+ - Calibration plots for training θ (PIT histogram + coverage curve).
+ * - ``tfs_summarize_theta_test_calibration.pdf``
+ - Calibration plots for test θ.
+ * - ``tfs_summarize_growth_calibration.pdf``
+ - Calibration plots for growth predictions.
+ * - ``tfs_summarize_params_{name}.pdf`` / ``.csv``
+ - Per-parameter scatter of simulated vs inferred values (simulated
+ data only).
+ * - ``tfs_summarize_params_{name}_calibration.pdf``
+ - Calibration plots for each inferred parameter (simulated data only).
+
+Reading the outputs
+-------------------
+
+.. _fit-summary-json:
+
+fit_summary.json — at a glance
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The JSON file gives a rapid numerical overview of the run. Open it in any
+text editor or load it with ``json.load``.
+
+.. code-block:: json
+
+ {
+ "metadata": {
+ "n_parameters": 5005,
+ "n_theta_training_points": 32,
+ "n_theta_test_points": 3752,
+ "n_growth_training_points": 54180,
+ "final_loss": 7281255.0
+ },
+ "theta": {
+ "training": { "pearson_r": 0.999, "r_squared": 0.998, "rmse": 0.016 },
+ "test": { "pearson_r": 0.936, "r_squared": 0.876, "rmse": 0.140 }
+ },
+ "growth": {
+ "training": { "pearson_r": 0.981, "r_squared": 0.963, "rmse": 0.577 }
+ }
+ }
+
+Key fields:
+
+* ``n_parameters`` — total number of free parameters (from the guesses CSV).
+ Compare to ``n_growth_training_points`` to check for over-parameterisation.
+* ``final_loss`` — the converged SVI ELBO (negative; more negative = better).
+ Use this to compare runs with identical data but different components or seeds.
+* ``theta.training`` — statistics comparing the model's θ predictions against
+ the direct **binding observations** used as training data. These will
+ always be near-perfect because the model is conditioned on these values.
+* ``theta.test`` — statistics comparing θ predictions against the full
+ ground-truth θ grid from ``tfs-simulate``. This is the real generalization
+ diagnostic: the vast majority of test points were never directly observed.
+* ``growth.training`` — statistics comparing predicted ln(CFU) against the
+ observed growth data.
+
+Useful derived checks:
+
+* ``theta.test.r_squared > 0.85`` is a reasonable threshold for a well-fitting
+ model on a library of this size and noise level.
+* A large gap between ``theta.training.rmse`` and ``theta.test.rmse`` signals
+ that the model fits the spiked genotypes well but struggles to generalize
+ to the broader library.
+* ``residual_corr`` and its p-value test for autocorrelation in the residuals.
+ A significant value indicates systematic structure the model is not capturing.
+
+Training-loss curve (tfs_summarize_losses.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The loss curve shows the SVI ELBO at each recorded epoch. A healthy run
+shows a smooth decrease that levels off to a plateau; a run that has not
+converged will still be declining at the final epoch. If the loss spikes
+upward mid-run, the learning rate or the prior scales may need adjustment.
+
+Theta correlation (tfs_summarize_theta_corr.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. figure:: _static/saa_theta_corr.png
+ :alt: Two-panel theta correlation plot
+ :width: 100%
+
+ **Left panel** — training θ: the model's posterior-median prediction
+ vs the direct binding observations used as input. Near-perfect agreement
+ is expected and is not by itself a measure of model quality. **Right
+ panel** — test θ: predictions vs the full ground-truth θ grid from
+ simulation, covering all library genotypes at all concentrations. Points
+ far from the diagonal indicate genotypes or concentrations where the
+ mutation-additivity assumption breaks down.
+
+The two panels serve different purposes:
+
+* **Training panel** confirms that the model correctly ingest the binding
+ data — any large deviation here points to a data-loading or configuration
+ error.
+* **Test panel** is the generalization diagnostic. The model must predict θ
+ for thousands of genotypes that were never directly measured; it does so by
+ combining per-mutation effects inferred from the spiked controls and the
+ growth data. Scatter around the diagonal reflects the irreducible noise in
+ that extrapolation.
+
+Per-genotype theta fits (tfs_summarize_{genotype}_theta_fits.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. figure:: _static/saa_wt_theta_fits.png
+ :alt: Wild-type theta fit
+ :width: 60%
+
+ Wild-type θ curve: posterior-median prediction (line) overlaid with the
+ binding observations (circles). The x-axis is on a log scale. A sigmoid
+ transition from high θ at low [IPTG] to low θ at high [IPTG] is the
+ expected shape for a repressor inactivated by its inducer.
+
+One plot is produced per genotype that appears in both the binding data and
+the theta predictions. For spiked genotypes with measured binding curves
+these plots confirm that the Hill model adequately describes the shape of the
+induction curve.
+
+Growth trajectories (tfs_summarize_{genotype}_trajectory.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The trajectory plots show predicted ln(CFU) versus time for every growth
+condition, with the observed data points (two replicates) overlaid on the
+posterior-median prediction and its 95% credible interval.
+
+One page is produced per genotype; each page contains one panel per
+condition (library × selection × titrant concentration). The dashed
+vertical line marks the transition from pre-growth to selection medium.
+
+Typical patterns to look for:
+
+* **Pre-growth panels** (``kanR-kan``, ``pheS-4CP`` without selection):
+ all genotypes should grow at similar rates; large deviations suggest a
+ problematic ``dk_geno`` estimate.
+* **Selection panels at extreme IPTG**: the two markers should respond in
+ opposite directions as θ moves from ~1 to ~0 with increasing [IPTG].
+* **Replicate agreement**: the two colored lines should be close together;
+ large replicate-to-replicate variation indicates high tube noise or a
+ confounded sample.
+
+Growth correlation (tfs_summarize_growth_corr.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. figure:: _static/saa_growth_corr.png
+ :alt: Growth prediction correlation
+ :width: 65%
+
+ Observed vs predicted ln(CFU) across all genotypes, conditions, and
+ time-points in the training data (54,180 points in the example run).
+ Points cluster tightly along the diagonal except at low ln(CFU) values
+ (~3–7), where sequencing noise dominates and the model's predictions
+ spread somewhat.
+
+A tight cluster along the diagonal confirms that the growth model
+accurately reproduces the data. Systematic vertical or horizontal bands
+point to conditions or time-points where the model consistently over- or
+under-predicts; these are usually caused by incorrect ``{m, b}`` priors for
+those conditions or by tube-noise events in specific samples.
+
+.. _calibration-plots:
+
+Calibration plots (tfs_summarize_*_calibration.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. figure:: _static/saa_growth_calibration.png
+ :alt: Growth calibration plots
+ :width: 100%
+
+ Growth-prediction calibration. **Left** — PIT histogram; the dashed
+ red line shows the ideal uniform distribution. **Right** — calibration
+ curve; the dashed black line is perfect calibration.
+
+Calibration plots assess whether the model's **posterior uncertainty
+intervals** are appropriately sized — neither too narrow (overconfident)
+nor too wide (underconfident).
+
+**PIT histogram** (left panel)
+
+The Probability Integral Transform (PIT) value for each observation is the
+fraction of the posterior predictive distribution that falls below the
+observed value. For a perfectly calibrated model, PIT values are
+uniformly distributed on [0, 1] and the histogram should be flat.
+
+* **U-shaped histogram** (spikes near 0 and 1, flat middle): many
+ observations fall outside the predicted intervals → posterior is
+ **overconfident** (intervals too narrow).
+* **Hump-shaped histogram** (peak in the middle, low edges): observations
+ cluster inside the predicted intervals → posterior is **underconfident**
+ (intervals too wide).
+
+**Calibration curve** (right panel)
+
+For each nominal coverage level *α* (x-axis), the calibration curve shows
+the empirical fraction of observations that fall within the model's *α*
+credible interval (y-axis).
+
+* **Below the diagonal**: empirical coverage < nominal → **overconfident**.
+* **Above the diagonal**: empirical coverage > nominal → **underconfident**.
+* **On the diagonal**: perfectly calibrated.
+
+The example growth-calibration plot shows U-shaped PIT and a curve below
+the diagonal. This slight overconfidence in growth predictions is typical
+of SVI: the mean-field variational approximation tends to underestimate
+posterior variance. It does not indicate a model misspecification; the
+point predictions (correlation, R²) remain accurate.
+
+Similar calibration plots are produced for theta (training and test) and for
+each parameter group when ground-truth values are available.
+
+Parameter recovery (tfs_summarize_params_{name}.pdf)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. figure:: _static/saa_params_log_hill_K.png
+ :alt: log_hill_K parameter recovery
+ :width: 65%
+
+ Simulated vs inferred ``log_hill_K`` (log₁₀ of the Hill dissociation
+ constant) for all library genotypes. Each point is one genotype; the
+ dashed line is perfect recovery. Good agreement across the dynamic range
+ confirms that the model correctly identifies which genotypes have shifted
+ binding affinities.
+
+Parameter-recovery plots are produced only when the run directory contains
+a ``*_sim_parameters.csv`` file (i.e. when running on simulated data). They
+compare the posterior-median inferred value for each parameter against the
+ground truth used during simulation.
+
+One plot is produced per parameter group. Versions are produced for:
+
+* **Direct parameters** — per-genotype values (e.g. ``log_hill_K``,
+ ``theta_high``, ``theta_low``, ``hill_n``).
+* **Diff parameters** (``_d_`` infix) — per-mutation effect sizes relative to
+ wild-type (e.g. ``d_log_hill_K``, ``d_logit_low``).
+* **Epistasis parameters** (``_epi_`` infix) — pairwise epistatic deviations
+ from additivity for double-mutant genotypes.
+
+Each parameter group also has a companion calibration plot
+(``*_calibration.pdf``) assessing whether the inferred uncertainty correctly
+covers the true value.
+
+Interpreting parameter recovery
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* Tight scatter around the diagonal for all parameters → model recovers the
+ ground truth; the experimental design has sufficient signal.
+* Tight recovery for direct parameters but wide scatter for diff/epi
+ parameters → the model identifies which genotypes differ from wild-type
+ but cannot resolve the per-mutation contributions as precisely.
+* Systematic bias (all points shifted to one side of the diagonal) →
+ mis-specified priors, an incorrect model component, or a calibration issue
+ in the pre-fit step.
+* Good recovery for ``log_hill_K`` but poor for ``hill_n`` or ``theta_high``
+ → the mid-range of the induction curve is well-constrained but the shape
+ of the saturation plateaux is not (common when the growth data covers only
+ a restricted θ range for most genotypes).
diff --git a/examples/simulate-and-analyze/README.md b/examples/simulate-and-analyze/README.md
new file mode 100644
index 00000000..c3765eb3
--- /dev/null
+++ b/examples/simulate-and-analyze/README.md
@@ -0,0 +1,33 @@
+# simulate-and-analyze
+
+A complete end-to-end example: simulate a synthetic TF-library screen and fit
+the hierarchical Bayesian model to the simulated data.
+
+## Contents
+
+| File | Purpose |
+|------|---------|
+| `simulate_config.yaml` | Simulation parameters (library genetics, conditions, binding data) |
+| `hill_params.csv` | Per-genotype Hill binding parameters used to generate ground-truth θ |
+| `run.sh` | Pipeline script (simulate → configure → calibrate → fit → predict → summarise) |
+
+## Running
+
+```bash
+bash run.sh simulate_config.yaml out/ 1
+```
+
+The three arguments are the config file, the output directory (created
+automatically), and the random seed. Seed `1` reproduces the documented
+example outputs.
+
+Running time is roughly 20–60 minutes on a laptop CPU or a few minutes on a
+GPU. On a cluster, comment out the `XLA_FLAGS` line in `run.sh` and
+uncomment `module load cuda/...` instead.
+
+## Documentation
+
+See the [quickstart](https://tfscreen.readthedocs.io/en/latest/quickstart.html)
+for a guided walkthrough and the
+[tfs-summarize-fit reference](https://tfscreen.readthedocs.io/en/latest/summarize-fit.html)
+for a detailed guide to interpreting the diagnostic outputs in `out/summary/`.
diff --git a/examples/simulate-and-analyze/hill_params.csv b/examples/simulate-and-analyze/hill_params.csv
new file mode 100644
index 00000000..a3568471
--- /dev/null
+++ b/examples/simulate-and-analyze/hill_params.csv
@@ -0,0 +1,5 @@
+genotype,log_hill_K,hill_n,theta_high,theta_low
+wt,-4.1,2,0.01,0.99
+M42I,-2.2104779610300698,1.5156149571654287,0.5054105302329879,0.9538933267956679
+H74A,-5.872705439673401,1.2001278450518955,0.05,0.7115380202706272
+K84L,-5.755115159475512,1.0593068151944411,0.1,0.856776050040223
diff --git a/examples/simulate-and-analyze/run.sh b/examples/simulate-and-analyze/run.sh
new file mode 100644
index 00000000..65d9b4fb
--- /dev/null
+++ b/examples/simulate-and-analyze/run.sh
@@ -0,0 +1,132 @@
+#!/bin/bash -l
+#SBATCH --account=harmslab ### change this to your actual account for charging
+#SBATCH --job-name=tfscreen ### job name
+#SBATCH --output=hostname.out ### file in which to store job stdout
+#SBATCH --error=hostname.err ### file in which to store job stderr
+#SBATCH --partition=gpu
+#SBATCH --time=01-00:00:00
+#SBATCH --nodes=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gpus=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --mem=0
+
+# On a cluster: comment out the XLA_FLAGS line and uncomment "module load cuda/..."
+# On a local CPU: keep the XLA_FLAGS line (sets JAX to use 8 virtual CPU devices).
+#module load cuda/12.4.1
+export XLA_FLAGS="--xla_force_host_platform_device_count=8"
+
+# Stop immediately on any error.
+set -e
+
+# ---------------------------------------------------------------------------
+# Parse arguments
+# ---------------------------------------------------------------------------
+config_file="${1}"
+run_dir="${2}"
+if [[ ! "${run_dir}" || ! "${config_file}" ]]; then
+ echo "Usage: run.sh config_file out_dir [seed]"
+ exit 1
+fi
+
+seed="${3}"
+if [[ ! "${seed}" ]]; then
+ seed=1
+fi
+
+# ---------------------------------------------------------------------------
+# 1. Simulate library
+# ---------------------------------------------------------------------------
+# tfs-simulate reads the config YAML and writes simulated growth data,
+# binding curves, and ground-truth parameter CSVs to run_dir.
+echo ">>> Simulate library"
+tfs-simulate "${config_file}" "${run_dir}" --seed "${seed}"
+
+cd "${run_dir}"
+
+# ---------------------------------------------------------------------------
+# 2. Configure model
+# ---------------------------------------------------------------------------
+# tfs-configure-model validates the data, maps labels to indices, selects
+# model components, and writes tfs_configure_config.yaml + priors/guesses CSVs.
+# Edit the flags here to change which model components are used.
+echo ">>> Configure model"
+tfs-configure-model \
+ tfs_sim_binding.csv \
+ --growth_df tfs_sim_growth.csv \
+ --presplit_df tfs_sim_presplit.csv \
+ --condition_growth_model linear \
+ --growth_transition_model instant \
+ --ln_cfu0_model hierarchical_factored \
+ --dk_geno_model hierarchical_geno \
+ --activity_model fixed \
+ --theta_model hill_mut \
+ --transformation_model single \
+ --theta_rescale_model passthrough \
+ --theta_growth_noise_model logit_normal \
+ --theta_binding_noise_model zero \
+ --growth_noise_model normal_kt \
+ --spiked wt M42I H74A K84L M42I/H74A M42I/K84L H74A/K84L D88A \
+ --growth_shares_replicates \
+ --epistasis
+
+# ---------------------------------------------------------------------------
+# 3. Pre-fit calibration
+# ---------------------------------------------------------------------------
+# A fast MAP fit on a simplified model calibrates the growth-linking-function
+# priors (m and b) before the full inference run. Updates priors.csv and
+# guesses.csv in place.
+echo ">>> Pre-fit calibration"
+tfs-prefit-calibration tfs_configure_config.yaml --seed "${seed}"
+
+# ---------------------------------------------------------------------------
+# 4. Fit model (SVI)
+# ---------------------------------------------------------------------------
+# Main hierarchical Bayesian inference. SVI produces a full approximate
+# posterior; for a point estimate only, use --analysis_method map instead.
+echo ">>> Fit model"
+tfs-fit-model \
+ tfs_configure_config.yaml \
+ --seed "${seed}" \
+ --analysis_method svi \
+ --convergence_tolerance 0.0001
+
+# ---------------------------------------------------------------------------
+# 5. Sample posterior
+# ---------------------------------------------------------------------------
+# Draw posterior samples from the SVI variational distribution and write
+# them to an HDF5 file used by the prediction steps below.
+echo ">>> Sample posterior"
+tfs-sample-posterior tfs_configure_config.yaml tfs_fit_model_checkpoint.pkl
+
+# ---------------------------------------------------------------------------
+# 6. Extract parameter estimates
+# ---------------------------------------------------------------------------
+# Summarise the posterior into per-parameter CSV files (quantiles, means).
+echo ">>> Extract parameter estimates"
+tfs-extract-params tfs_configure_config.yaml tfs_posterior.h5
+
+# ---------------------------------------------------------------------------
+# 7. Predict theta
+# ---------------------------------------------------------------------------
+# Predict operator occupancy θ as a function of titrant concentration for
+# every genotype in the training data.
+echo ">>> Predict theta"
+tfs-predict-theta tfs_configure_config.yaml tfs_posterior.h5
+
+# ---------------------------------------------------------------------------
+# 8. Predict growth
+# ---------------------------------------------------------------------------
+# Predict ln(CFU) with posterior uncertainty for every training observation.
+echo ">>> Predict growth"
+tfs-predict-growth tfs_configure_config.yaml tfs_posterior.h5 --num_marginal_samples=500
+
+# ---------------------------------------------------------------------------
+# 9. Summarise fit
+# ---------------------------------------------------------------------------
+# Collects all outputs, computes statistics, and writes diagnostic PDFs and
+# CSVs to the summary/ subdirectory.
+echo ">>> Summarise fit"
+tfs-summarize-fit .
+
+cd ..
diff --git a/examples/simulate-and-analyze/simulate_config.yaml b/examples/simulate-and-analyze/simulate_config.yaml
new file mode 100644
index 00000000..a0ea8afa
--- /dev/null
+++ b/examples/simulate-and-analyze/simulate_config.yaml
@@ -0,0 +1,246 @@
+
+# ---------------------------------------------------------------------------
+# Library genetic information
+# ---------------------------------------------------------------------------
+
+# Reading frame 0, 1, 2. No reverse (negative) strand allowed.
+reading_frame: 0
+
+# Amino acid residue number for the first in-frame residue of the amplicon.
+first_amplicon_residue: 27
+
+# Wildtype DNA sequence of the amplicon.
+wt_seq: gccagccacgtttctgcgaaaacgcgggaaaaagtggaagcggcgatggcggagctgaattacattcccaaccgcgtggcacaacaactggcgggcaaacagtcgttgctgattggcgttgccacctccagtctggccctgcacgcgccgtcgcaaattgtcgcggcgattaaatctcgcgccgatcaactgggtgccagcgtggtggtgtcgatg
+
+# Degenerate codon pattern for library positions. '.' = wildtype; 'nnt' = NNT
+# codon (encodes all 20 amino acids + stop with reduced synonymous redundancy).
+# Must be the same length as wt_seq.
+degen_sites: |
+ .............................................nnt.............................................................................................nnt...........................nnt..........................................
+# ......nntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnnt......nntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnntnnt
+
+
+# sub_libraries groups contiguous degenerate positions that were cloned together.
+# '.' = wildtype; any other character defines a sub-library block. Blocks must
+# be contiguous (......111... ok; ....11.1... not ok).
+sub_libraries: |
+ ......111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111......222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222222
+
+# Flanking sequences immediately upstream and downstream of the amplicon.
+expected_5p: cgcgtggtgaaccag
+expected_3p: gtagaacgaagcggc
+
+# Which genotype classes to include in the library. 'single-x' = all
+# single-mutation variants in sub-library x. 'double-x-y' = all pairwise
+# combinations between sub-libraries x and y. 'double-x-x' = internal
+# doubles within sub-library x.
+library_combos: ["single-1", "single-2", "double-1-2"]
+
+# Control sequences (e.g. known variants) added as spikes. Must be the same
+# length as wt_seq.
+spiked_seqs:
+- ........................................................................................................................................................................................................................ # wt
+- .............................................ata........................................................................................................................................................................ # m42i
+- .............................................................................................................................................gcg........................................................................ # h74a
+- ...........................................................................................................................................................................ctg.......................................... # k84l
+- .............................................ata.............................................................................................gcg........................................................................ # m42i/h74a
+- .............................................ata...........................................................................................................................ctg.......................................... # m42i/k84l
+- .............................................................................................................................................gcg...........................ctg.......................................... # h74a/k84l
+- .............................................ata.............................................................................................gcg...........................ctg.......................................... # m42i/h74a/k84l
+- ........................................................................................................................................................................................cc.............................. # d88a
+
+# ---------------------------------------------------------------------------
+# Phenotype calculation
+# ---------------------------------------------------------------------------
+
+# Registered theta component name. Any key in model_registry["theta"] is
+# valid except "_simple" (calibration-only). Examples: "hill_geno", "thermo.O2_C12_K5_U0_a.PK".
+theta_component: "hill_mut"
+
+# Pairwise epistasis (regularized horseshoe; used when library has doubles):
+theta_sim_priors:
+ epi_tau_scale: 0.1 # HalfCauchy scale for global sparsity τ
+
+# Optional: path to structural/thermodynamic data file required by PnnC and
+# PddG theta components. Omit (or set to null) for all other models.
+# .h5 / .hdf5 → HDF5 structural ensemble (PnnC)
+# .csv → per-mutation per-structure ddG prior means (PddG)
+# thermo_data: null
+
+# Optional: overrides for the theta component's default hyperparameters.
+# Any key returned by the component's get_hyperparameters() can be overridden.
+# Omit this block to use all component defaults.
+# theta_priors:
+# theta_tf_total_M: 3.3e-7
+
+# ---------------------------------------------------------------------------
+# Conditions
+# ---------------------------------------------------------------------------
+
+# Each block defines one set of growth conditions. The Cartesian product of
+# titrant_conc × t_sel gives one sample per entry. condition_pre and
+# condition_sel must also appear as keys in the 'growth' section below.
+condition_blocks:
+ - library: "kanR"
+ titrant_name: "iptg"
+ titrant_conc: [0, 0.0001, 0.001, 0.003]
+ condition_pre: "kanR-kan"
+ t_pre: 30
+ condition_sel: "kanR+kan"
+ t_sel: [160, 180, 200]
+
+ - library: "kanR"
+ titrant_name: "iptg"
+ titrant_conc: [0.01, 0.03, 0.1, 1.0]
+ condition_pre: "kanR-kan"
+ t_pre: 30
+ condition_sel: "kanR+kan"
+ t_sel: [145, 165, 180]
+
+ - library: "kanR"
+ titrant_name: "iptg"
+ titrant_conc: [0, 1.0]
+ condition_pre: "kanR-kan"
+ t_pre: 30
+ condition_sel: "kanR-kan"
+ t_sel: [160, 180, 200]
+
+ - library: "pheS"
+ titrant_name: "iptg"
+ titrant_conc: [0, 0.0001, 0.001, 0.003, 0.01, 0.03, 0.1, 1.0]
+ condition_pre: "pheS-4CP"
+ t_pre: 30
+ condition_sel: "pheS+4CP"
+ t_sel: [95, 110, 125]
+
+ - library: "pheS"
+ titrant_name: "iptg"
+ titrant_conc: [0, 1.0]
+ condition_pre: "pheS-4CP"
+ t_pre: 30
+ condition_sel: "pheS-4CP"
+ t_sel: [95, 110, 125]
+
+# ---------------------------------------------------------------------------
+# Growth rate parameters
+# ---------------------------------------------------------------------------
+
+# Linear mapping from operator occupancy (theta) to growth rate k for each
+# condition: k = m * theta + b. Every condition_pre and condition_sel that
+# appears in condition_blocks must have an entry here, and vice versa.
+growth:
+ kanR+kan: {m: -0.009933, b: 0.010696}
+ kanR-kan: {m: 0.000808, b: 0.015366}
+ pheS+4CP: {m: 0.006226, b: 0.021437}
+ pheS-4CP: {m: -0.000344, b: 0.028558}
+
+# Per-genotype pleiotropic growth cost is sampled from a shifted lognormal:
+# dk_geno = hyper_shift - exp(Normal(hyper_loc, hyper_scale))
+# Most variants are mildly deleterious; a short right tail (at most hyper_shift)
+# allows beneficial mutations. Matches the hierarchical tfmodel inference prior.
+dk_geno_hyper_loc: -3.5
+dk_geno_hyper_scale: 0.5
+dk_geno_hyper_shift: 0.02
+dk_geno_zero: false
+
+# Per-genotype TF activity (A) — a scalar that multiplies the theta contribution
+# to growth: k = b + A * m * theta. The wild-type is fixed at activity_wt=1.0
+# by inference convention. Set activity_mut_scale > 0 to draw each mutant
+# genotype's log(A) independently from Normal(log(activity_wt), activity_mut_scale).
+# activity_mut_scale=0 (default) gives all genotypes identical activity, matching
+# the 'fixed' activity inference component.
+activity_wt: 1.0
+activity_mut_scale: 0.0
+
+# ---------------------------------------------------------------------------
+# Experimental simulation parameters
+# ---------------------------------------------------------------------------
+
+# Number of transformants per sub-library. Keys must match the characters in
+# library_combos (e.g. 'single-1', 'double-1-2') plus 'spiked'.
+transform_sizes:
+ single-1: 100_000
+ single-2: 100_000
+ double-1-2: 300_000
+ spiked: 1000
+
+# Mixing ratio of sub-libraries in the pooled library. Only the ratio matters.
+library_mixture:
+ single-1: 100
+ single-2: 100
+ double-1-2: 1000
+ spiked: 100
+
+# Sigma of the log-normal used to draw initial genotype frequencies. 0 = flat
+# distribution; 1–2 is realistic for a typical transformation.
+lib_assembly_skew_sigma: 1.5
+
+# Each bacterium receives a number of plasmids drawn from a zero-truncated
+# Poisson with this lambda. Set to 0 or null for exactly one plasmid per cell.
+transformation_poisson_lambda: 0
+
+# When a cell carries multiple plasmids, growth rates are combined via this
+# function. Recommended: 'min' for repressors of antibiotic resistance (tightest
+# binder dominates); 'max' for inducers. Default: 'mean' for all libraries.
+multi_plasmid_combine_fcn:
+ kanR: "min"
+ pheS: "max"
+
+# Colony-forming units per mL at the start of pre-growth.
+cfu0: 6.5e7
+
+# Per-tube environmental growth-rate noise (hr⁻¹). Each physical tube gets its
+# own delta_k ~ Normal(0, tube_noise_sigma); the kt contribution per tube is
+# delta_k * (t_pre + t_sel), shared by all genotypes in that tube. A value of
+# ~0.001–0.002 hr⁻¹ is typical for well-controlled liquid cultures.
+# Set to 0 or null to disable.
+tube_noise_sigma: 0.002
+
+# ---------------------------------------------------------------------------
+# Growth transition model (optional)
+# ---------------------------------------------------------------------------
+
+# Models the lag phase when bacteria switch from pre-selection to selection medium.
+# Every condition_pre in condition_blocks must have an entry. Use model: "instant"
+# for conditions with no detectable lag. Omit this block to use instant everywhere.
+#
+# "memory" model: tau = tau0 + k1 / (theta + k2)
+# tau0 = baseline lag (minutes)
+# k1 = occupancy-dependence of lag (higher = more effect)
+# k2 = curvature constant (> 0; prevents division by zero)
+growth_transition:
+ - condition_pre: "pheS-4CP"
+ model: "instant"
+
+ - condition_pre: "kanR-kan"
+ model: "instant"
+
+# ---------------------------------------------------------------------------
+# Data collection parameters
+# ---------------------------------------------------------------------------
+
+# Total sequencing reads across all samples.
+total_num_reads: 100_000_000
+
+# Fraction of reads re-assigned to a randomly chosen genotype (index hopping).
+# Set to 0 or null to disable.
+prob_index_hop: 0.00
+
+# Random seed for reproducibility. Set to null for a random seed each run.
+seed: 12345
+
+# ---------------------------------------------------------------------------
+# Binding data simulation (optional)
+# ---------------------------------------------------------------------------
+
+# If present, tfs-simulate also writes a binding CSV with per-genotype theta
+# curves. titrant_name must match e_name in observable_calc_kwargs.
+
+binding_data:
+ genotype_params_file: hill_params.csv
+ titrant_name: iptg
+ titrant_conc: [0, 0.0001, 0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1.0]
+ noise: 0.001
+
+presplit_data:
+ noise: 0.00
diff --git a/examples/simulate/simulate_config.yaml b/examples/simulate/simulate_config.yaml
index cde1b22a..33362a67 100644
--- a/examples/simulate/simulate_config.yaml
+++ b/examples/simulate/simulate_config.yaml
@@ -43,7 +43,29 @@ spiked_seqs:
# ---------------------------------------------------------------------------
# Registered theta component name. Any key in model_registry["theta"] is
-# valid except "_simple" (calibration-only). Examples: "hill_geno", "thermo.O2_C12_K5_U0_a.PK".
+# valid except "_simple" (calibration-only).
+#
+# Thermodynamic components (physics-based operator-occupancy model):
+# "thermo.O2_C12_K5_U0_a.PK" — dimer with operator architecture O2, 12-mer
+# cooperative lattice, 5 binding sites, etc.
+# (other thermo.* keys available — see model_registry)
+#
+# Perturbation-based components (wildtype reference + per-genotype perturbations):
+# "hill_geno" — per-genotype Hill equation; each genotype's (theta_low,
+# theta_high, K, n) drawn independently around a WT reference.
+# Supports four phenotype categories: normal, stuck-bound,
+# never-binds, inverted-logic.
+# "hill_mut" — mutation-decomposed Hill equation; per-mutation additive
+# deltas in transformed parameter space, assembled via
+# matrix multiplication. Pairwise epistasis sampled from a
+# regularized horseshoe prior when the library includes
+# double mutants. Wildtype genotype always receives the
+# exact reference curve.
+#
+# Perturbation-based components ignore theta_priors and instead use
+# theta_sim_priors (see below). They do not call define_model; tfs-simulate
+# will fall back to prior-predictive sampling only for components that lack a
+# simulate() function.
theta_component: "thermo.O2_C12_K5_U0_a.PK"
# Optional: path to structural/thermodynamic data file required by PnnC and
@@ -52,14 +74,60 @@ theta_component: "thermo.O2_C12_K5_U0_a.PK"
# .csv → per-mutation per-structure ddG prior means (PddG)
# thermo_data: null
-# Optional: overrides for the theta component's default hyperparameters.
+# Optional: overrides for the theta component's default model hyperparameters.
# Any key returned by the component's get_hyperparameters() can be overridden.
+# Used by thermodynamic/hierarchical components; ignored when theta_component
+# is "hill_geno" or "hill_mut" (those use theta_sim_priors instead).
# Omit this block to use all component defaults.
# theta_priors:
# theta_tf_total_M: 3.3e-7
-# Integer seed for the JAX RNG used when sampling theta from the prior.
-theta_rng_seed: 0
+# Optional: overrides for perturbation-based simulation priors (SimPriors).
+# Used only when theta_component is "hill_geno" or "hill_mut".
+# Any key returned by the component's get_sim_hyperparameters() can be overridden.
+# Omit this block to use all component defaults.
+#
+# --- hill_geno SimPriors keys ---
+# Wildtype reference curve (a full-range decreasing Hill curve by default):
+# wt_theta_low: 0.99 # theta at zero ligand (near fully bound)
+# wt_theta_high: 0.01 # theta at saturating ligand (near unbound)
+# wt_log_K: -4.1 # ln(K_D in mM); sets midpoint (lac/IPTG default)
+# wt_hill_n: 2.0 # Hill coefficient (cooperativity)
+# Per-genotype Normal perturbation widths (in logit/log space):
+# sigma_logit_low: 0.5 # spread on logit(theta_low)
+# sigma_logit_delta: 0.5 # spread on logit_delta = logit_high − logit_low
+# sigma_log_K: 0.5 # spread on log(K) — primary binding-affinity effect
+# sigma_log_n: 0.3 # spread on log(hill_n)
+# Mixture fractions for special phenotype categories (must sum to < 1;
+# remainder is "normal"):
+# p_stuck_bound: 0.05 # fraction whose theta ≈ 1 everywhere
+# p_never_binds: 0.05 # fraction whose theta ≈ 0 everywhere
+# p_inverted: 0.02 # fraction with inverted logic (theta rises with ligand)
+#
+# --- hill_mut SimPriors keys ---
+# Wildtype reference (shared with hill_geno defaults):
+# wt_theta_low: 0.99
+# wt_theta_high: 0.01
+# wt_log_K: -4.1
+# wt_hill_n: 2.0
+# Per-mutation delta widths (additive in logit/log space):
+# sigma_d_logit_low: 0.3
+# sigma_d_logit_delta: 0.5
+# sigma_d_log_K: 0.5 # largest effect — primarily shifts binding affinity
+# sigma_d_log_n: 0.3
+# Pairwise epistasis (regularized horseshoe; used when library has doubles):
+# epi_tau_scale: 0.1 # HalfCauchy scale for global sparsity τ
+# # set to 0.0 for exactly no epistasis
+# epi_slab_scale: 2.0 # typical large-epistasis effect size
+# epi_slab_df: 4.0 # InvGamma slab degrees of freedom
+#
+# Example — stress-test epistasis recovery:
+# theta_sim_priors:
+# epi_tau_scale: 1.0 # strong, dense epistasis
+#
+# Example — simulate with no epistasis at all:
+# theta_sim_priors:
+# epi_tau_scale: 0.0
# ---------------------------------------------------------------------------
# Conditions
@@ -113,6 +181,9 @@ growth:
dk_geno_hyper_loc: -3.5
dk_geno_hyper_scale: 1.0
dk_geno_hyper_shift: 0.02
+# Set dk_geno_zero: true to pin every genotype's dk_geno to 0, skipping the
+# stochastic draw above. When set, the three dk_geno_hyper_* keys are optional.
+# dk_geno_zero: true
# Per-genotype TF activity (A) — a scalar that multiplies the theta contribution
# to growth: k = b + A * m * theta. The wild-type is fixed at activity_wt=1.0
@@ -204,7 +275,7 @@ total_num_reads: 25_000_000
prob_index_hop: 0.05
# Random seed for reproducibility. Set to null for a random seed each run.
-random_seed: null
+seed: null
# ---------------------------------------------------------------------------
# Binding data simulation (optional)
@@ -220,3 +291,34 @@ random_seed: null
# titrant_name: iptg
# titrant_conc: [0, 0.0001, 0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1.0]
# noise: 0.02
+
+# ---------------------------------------------------------------------------
+# Pre-split data simulation (optional)
+# ---------------------------------------------------------------------------
+
+# tfs-simulate writes a presplit CSV when a 'presplit_data' key is present
+# in this file. To enable it, ADD the following lines (uncommented) anywhere
+# at the top level of this config — for example, right here:
+#
+# presplit_data:
+# noise: 0.0
+#
+# The 'noise' sub-key is optional; omitting it (or setting noise: 0.0) uses
+# sequencing-count uncertainty only. The minimal activation is just:
+#
+# presplit_data:
+#
+# ---- What the presplit CSV represents -------------------------------------
+# This simulates the sequencing aliquot taken at t = -t_pre, just before
+# the culture is split into titrant-concentration conditions. Because it is
+# taken before any pre-selection growth, each ln_cfu value directly
+# constrains ln_cfu0 in the hierarchical model.
+#
+# Output columns: replicate, condition_pre, genotype, ln_cfu, ln_cfu_std,
+# ln_cfu_0_true (ground truth for validation).
+# One row per (replicate, condition_pre, genotype).
+#
+# Pass the resulting file to tfs-configure-model with --presplit_df.
+#
+# noise (float, default 0): extra Gaussian noise on the ln-scale, added on
+# top of sequencing-count uncertainty. Useful for benchmarking.
diff --git a/examples/simulate/simulate_grid.yaml b/examples/simulate/simulate_grid.yaml
index 24a7cd74..a842023b 100644
--- a/examples/simulate/simulate_grid.yaml
+++ b/examples/simulate/simulate_grid.yaml
@@ -17,7 +17,7 @@ base_config: simulate_config.yaml # in the same examples/simulate/ directory
# Directory name for each run. Variables from both 'simulate' and 'template'
# blocks are available. The |basename filter strips directory paths from
# file-valued variables.
-run_name: "{{ theta_component }}__noise{{ tube_noise_sigma }}__seed{{ random_seed }}"
+run_name: "{{ theta_component }}__noise{{ tube_noise_sigma }}__seed{{ seed }}"
# Jinja2 template file (relative to this YAML) rendered into each subdir.
output_file: run.sh # in the same examples/simulate/ directory
@@ -65,8 +65,8 @@ simulate:
- name: seed
variants:
- - random_seed: 0
- - random_seed: 42
+ - seed: 0
+ - seed: 42
# ---------------------------------------------------------------------------
diff --git a/pyproject.toml b/pyproject.toml
index 02d47bae..7c12e73e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -73,16 +73,16 @@ docs = ["pydata-sphinx-theme"]
tfs-process-fastq = "tfscreen.process_raw.scripts.process_fastq_cli:main"
tfs-process-counts = "tfscreen.process_raw.scripts.process_counts_cli:main"
tfs-configure-model = "tfscreen.tfmodel.scripts.configure_model_cli:main"
-tfs-prefit-calibration = "tfscreen.tfmodel.scripts.run_prefit_calibration_cli:main"
+tfs-prefit-calibration = "tfscreen.tfmodel.scripts.prefit_calibration_cli:main"
tfs-fit-model = "tfscreen.tfmodel.scripts.fit_model_cli:main"
tfs-sample-posterior = "tfscreen.tfmodel.scripts.sample_posterior_cli:main"
-tfs-sample-prior = "tfscreen.tfmodel.scripts.prior_predictive_cli:main"
+tfs-sample-prior = "tfscreen.tfmodel.scripts.sample_prior_cli:main"
tfs-extract-params = "tfscreen.tfmodel.scripts.extract_params_cli:main"
tfs-predict-growth = "tfscreen.tfmodel.scripts.predict_growth_cli:main"
tfs-predict-theta = "tfscreen.tfmodel.scripts.predict_theta_cli:main"
-tfs-cat-response = "tfscreen.analysis.cat_response.cat_response_cli:main"
+tfs-cat-response = "tfscreen.analysis.cat_response.scripts.cat_response_cli:main"
tfs-diagnose-nan = "tfscreen.tfmodel.scripts.diagnose_nan_cli:main"
-tfs-simulate = "tfscreen.simulate.scripts.run_simulation_cli:main"
+tfs-simulate = "tfscreen.simulate.scripts.simulate_cli:main"
tfs-setup-sim-grid = "tfscreen.simulate.scripts.setup_sim_grid_cli:main"
tfs-setup-grid = "tfscreen.tfmodel.scripts.setup_grid_cli:main"
tfs-summarize-grid = "tfscreen.tfmodel.scripts.summarize_grid_cli:main"
diff --git a/reports/flake.txt b/reports/flake.txt
index 55c3da9b..b653ba9c 100644
--- a/reports/flake.txt
+++ b/reports/flake.txt
@@ -3,7 +3,7 @@
./scripts/generate_struct_ensemble.py:433:18: E221 multiple spaces before operator
./src/tfscreen/__init__.py:19:1: W391 blank line at end of file
./src/tfscreen/__version__.py:5:16: W291 trailing whitespace
-./src/tfscreen/analysis/__init__.py:9:1: W391 blank line at end of file
+./src/tfscreen/analysis/__init__.py:13:1: W391 blank line at end of file
./src/tfscreen/analysis/cat_response/__init__.py:2:77: W291 trailing whitespace
./src/tfscreen/analysis/cat_response/__init__.py:3:55: W291 trailing whitespace
./src/tfscreen/analysis/cat_response/cat_fit.py:14:1: C901 'cat_fit' is too complex (19)
@@ -70,11 +70,9 @@
./src/tfscreen/analysis/cat_response/cat_response.py:134:34: E231 missing whitespace after ','
./src/tfscreen/analysis/cat_response/cat_response.py:138:42: E231 missing whitespace after ','
./src/tfscreen/analysis/cat_response/cat_response.py:141:56: W292 no newline at end of file
-./src/tfscreen/analysis/cat_response/cat_response_cli.py:12:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/analysis/cat_response/cat_response_cli.py:21:1: C901 'cat_response' is too complex (11)
-./src/tfscreen/analysis/cat_response/cat_response_cli.py:91:56: E127 continuation line over-indented for visual indent
-./src/tfscreen/analysis/cat_response/fit_response_cli.py:25:1: C901 'fit_response' is too complex (11)
-./src/tfscreen/analysis/cat_response/fit_response_cli.py:87:56: E127 continuation line over-indented for visual indent
+./src/tfscreen/analysis/cat_response/scripts/cat_response_cli.py:12:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/analysis/cat_response/scripts/cat_response_cli.py:21:1: C901 'cat_response' is too complex (11)
+./src/tfscreen/analysis/cat_response/scripts/cat_response_cli.py:91:56: E127 continuation line over-indented for visual indent
./src/tfscreen/analysis/extract_epistasis.py:8:1: C901 'mutant_cycle_pivot' is too complex (11)
./src/tfscreen/analysis/extract_epistasis.py:8:1: E302 expected 2 blank lines, found 1
./src/tfscreen/analysis/extract_epistasis.py:33:77: W291 trailing whitespace
@@ -103,6 +101,28 @@
./src/tfscreen/analysis/extract_epistasis.py:229:1: W293 blank line contains whitespace
./src/tfscreen/analysis/extract_epistasis.py:231:10: W291 trailing whitespace
./src/tfscreen/analysis/extract_epistasis.py:233:30: W291 trailing whitespace
+./src/tfscreen/analysis/stats_test_suite.py:8:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/analysis/stats_test_suite.py:8:31: E231 missing whitespace after ','
+./src/tfscreen/analysis/stats_test_suite.py:8:42: E231 missing whitespace after ','
+./src/tfscreen/analysis/stats_test_suite.py:55:75: W291 trailing whitespace
+./src/tfscreen/analysis/stats_test_suite.py:66:52: W291 trailing whitespace
+./src/tfscreen/analysis/stats_test_suite.py:81:1: W293 blank line contains whitespace
+./src/tfscreen/analysis/stats_test_suite.py:91:1: W293 blank line contains whitespace
+./src/tfscreen/analysis/stats_test_suite.py:104:33: W291 trailing whitespace
+./src/tfscreen/analysis/stats_test_suite.py:134:1: W293 blank line contains whitespace
+./src/tfscreen/analysis/stats_test_suite.py:143:26: E231 missing whitespace after ','
+./src/tfscreen/analysis/stats_test_suite.py:144:1: W293 blank line contains whitespace
+./src/tfscreen/analysis/stats_test_suite.py:148:5: E303 too many blank lines (2)
+./src/tfscreen/analysis/stats_test_suite.py:149:22: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:150:15: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:151:26: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:155:21: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:156:24: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:157:24: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:158:32: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:159:21: E231 missing whitespace after ':'
+./src/tfscreen/analysis/stats_test_suite.py:161:1: W293 blank line contains whitespace
+./src/tfscreen/analysis/stats_test_suite.py:163:1: W391 blank line at end of file
./src/tfscreen/genetics/build_cycles.py:6:1: E302 expected 2 blank lines, found 1
./src/tfscreen/genetics/build_cycles.py:26:1: W293 blank line contains whitespace
./src/tfscreen/genetics/build_cycles.py:69:26: W291 trailing whitespace
@@ -114,13 +134,13 @@
./src/tfscreen/genetics/combine_mutation_effects.py:34:71: W291 trailing whitespace
./src/tfscreen/genetics/combine_mutation_effects.py:35:21: W291 trailing whitespace
./src/tfscreen/genetics/combine_mutation_effects.py:65:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/combine_mutation_effects.py:72:27: E261 at least two spaces before inline comment
-./src/tfscreen/genetics/combine_mutation_effects.py:73:10: E111 indentation is not a multiple of 4
-./src/tfscreen/genetics/combine_mutation_effects.py:73:10: E117 over-indented
-./src/tfscreen/genetics/combine_mutation_effects.py:81:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/combine_mutation_effects.py:88:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/combine_mutation_effects.py:102:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/combine_mutation_effects.py:103:25: W292 no newline at end of file
+./src/tfscreen/genetics/combine_mutation_effects.py:73:27: E261 at least two spaces before inline comment
+./src/tfscreen/genetics/combine_mutation_effects.py:74:10: E111 indentation is not a multiple of 4
+./src/tfscreen/genetics/combine_mutation_effects.py:74:10: E117 over-indented
+./src/tfscreen/genetics/combine_mutation_effects.py:82:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/combine_mutation_effects.py:89:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/combine_mutation_effects.py:103:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/combine_mutation_effects.py:104:25: W292 no newline at end of file
./src/tfscreen/genetics/count_mutation_backgrounds.py:8:1: C901 'count_mutation_backgrounds' is too complex (12)
./src/tfscreen/genetics/data.py:25:8: E231 missing whitespace after ':'
./src/tfscreen/genetics/data.py:26:8: E231 missing whitespace after ':'
@@ -228,10 +248,8 @@
./src/tfscreen/genetics/library_manager.py:251:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:254:27: E231 missing whitespace after ','
./src/tfscreen/genetics/library_manager.py:254:36: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:255:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:258:32: E231 missing whitespace after ','
./src/tfscreen/genetics/library_manager.py:258:46: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:259:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:264:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:268:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:269:82: W291 trailing whitespace
@@ -268,78 +286,78 @@
./src/tfscreen/genetics/library_manager.py:386:5: E303 too many blank lines (2)
./src/tfscreen/genetics/library_manager.py:415:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:421:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:425:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:464:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:465:78: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:469:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:474:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:476:25: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:476:38: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:477:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:480:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:486:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:491:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:524:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:525:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:526:5: E303 too many blank lines (2)
-./src/tfscreen/genetics/library_manager.py:553:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:558:73: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:558:82: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:559:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:560:41: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:560:55: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:560:63: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:561:39: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:561:43: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:562:34: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:562:38: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:564:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:427:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:466:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:467:78: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:471:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:476:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:478:25: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:478:38: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:479:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:482:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:488:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:493:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:526:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:527:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:528:5: E303 too many blank lines (2)
+./src/tfscreen/genetics/library_manager.py:555:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:560:73: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:560:82: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:561:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:562:41: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:562:55: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:562:63: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:563:39: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:563:43: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:564:34: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:564:38: E231 missing whitespace after ','
./src/tfscreen/genetics/library_manager.py:566:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:603:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:609:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:568:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:605:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:611:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:614:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:615:5: E303 too many blank lines (2)
-./src/tfscreen/genetics/library_manager.py:646:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:613:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:616:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:617:5: E303 too many blank lines (2)
./src/tfscreen/genetics/library_manager.py:648:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:654:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:655:31: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:659:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:650:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:656:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:657:31: E231 missing whitespace after ','
./src/tfscreen/genetics/library_manager.py:661:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:663:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:693:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:696:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:701:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:706:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:665:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:695:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:698:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:703:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:708:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:710:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:711:30: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:719:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:737:28: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:737:41: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:743:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:771:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:712:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:713:30: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:721:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:740:28: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:740:41: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:748:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:776:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:782:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:785:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:786:39: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:789:72: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:781:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:787:1: W293 blank line contains whitespace
./src/tfscreen/genetics/library_manager.py:790:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:793:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:801:9: E303 too many blank lines (2)
-./src/tfscreen/genetics/library_manager.py:802:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:805:76: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:806:15: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:820:1: W293 blank line contains whitespace
-./src/tfscreen/genetics/library_manager.py:833:79: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:834:75: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:835:34: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:836:58: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:840:64: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:848:37: W291 trailing whitespace
-./src/tfscreen/genetics/library_manager.py:849:27: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:851:40: E231 missing whitespace after ','
-./src/tfscreen/genetics/library_manager.py:854:1: W391 blank line at end of file
+./src/tfscreen/genetics/library_manager.py:791:39: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:794:72: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:795:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:798:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:806:9: E303 too many blank lines (2)
+./src/tfscreen/genetics/library_manager.py:807:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:810:76: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:811:15: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:825:1: W293 blank line contains whitespace
+./src/tfscreen/genetics/library_manager.py:838:79: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:839:75: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:840:34: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:841:58: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:845:64: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:853:37: W291 trailing whitespace
+./src/tfscreen/genetics/library_manager.py:854:27: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:856:40: E231 missing whitespace after ','
+./src/tfscreen/genetics/library_manager.py:859:1: W391 blank line at end of file
./src/tfscreen/mle/curve_models/__init__.py:4:77: W291 trailing whitespace
./src/tfscreen/mle/curve_models/__init__.py:5:42: W291 trailing whitespace
./src/tfscreen/mle/curve_models/__init__.py:11:77: W291 trailing whitespace
@@ -599,28 +617,6 @@
./src/tfscreen/mle/predict_with_error.py:57:43: E231 missing whitespace after ','
./src/tfscreen/mle/predict_with_error.py:61:45: E231 missing whitespace after ','
./src/tfscreen/mle/predict_with_error.py:74:32: W292 no newline at end of file
-./src/tfscreen/mle/stats_test_suite.py:8:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/mle/stats_test_suite.py:8:31: E231 missing whitespace after ','
-./src/tfscreen/mle/stats_test_suite.py:8:42: E231 missing whitespace after ','
-./src/tfscreen/mle/stats_test_suite.py:55:75: W291 trailing whitespace
-./src/tfscreen/mle/stats_test_suite.py:66:52: W291 trailing whitespace
-./src/tfscreen/mle/stats_test_suite.py:81:1: W293 blank line contains whitespace
-./src/tfscreen/mle/stats_test_suite.py:91:1: W293 blank line contains whitespace
-./src/tfscreen/mle/stats_test_suite.py:104:33: W291 trailing whitespace
-./src/tfscreen/mle/stats_test_suite.py:134:1: W293 blank line contains whitespace
-./src/tfscreen/mle/stats_test_suite.py:143:26: E231 missing whitespace after ','
-./src/tfscreen/mle/stats_test_suite.py:144:1: W293 blank line contains whitespace
-./src/tfscreen/mle/stats_test_suite.py:148:5: E303 too many blank lines (2)
-./src/tfscreen/mle/stats_test_suite.py:149:22: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:150:15: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:151:26: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:155:21: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:156:24: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:157:24: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:158:32: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:159:21: E231 missing whitespace after ':'
-./src/tfscreen/mle/stats_test_suite.py:161:1: W293 blank line contains whitespace
-./src/tfscreen/mle/stats_test_suite.py:163:1: W391 blank line at end of file
./src/tfscreen/plot/__init__.py:36:63: W291 trailing whitespace
./src/tfscreen/plot/cat_fits.py:13:1: C901 'cat_fits' is too complex (11)
./src/tfscreen/plot/cat_fits.py:13:1: E302 expected 2 blank lines, found 1
@@ -780,6 +776,7 @@
./src/tfscreen/plot/est_v_real_summary.py:102:1: W293 blank line contains whitespace
./src/tfscreen/plot/est_v_real_summary.py:104:1: W293 blank line contains whitespace
./src/tfscreen/plot/est_v_real_summary.py:106:1: W391 blank line at end of file
+./src/tfscreen/plot/geno_trajectory.py:29:1: C901 'plot_geno_trajectory' is too complex (14)
./src/tfscreen/plot/heatmap/aa_v_res_heatmap.py:6:1: E302 expected 2 blank lines, found 1
./src/tfscreen/plot/heatmap/aa_v_res_heatmap.py:15:53: E231 missing whitespace after ':'
./src/tfscreen/plot/heatmap/aa_v_res_heatmap.py:16:53: E231 missing whitespace after ':'
@@ -949,14 +946,13 @@
./src/tfscreen/plot/plot_theta_fits.py:108:46: E231 missing whitespace after ','
./src/tfscreen/plot/plot_theta_fits.py:115:42: E231 missing whitespace after ','
./src/tfscreen/plot/plot_theta_fits.py:121:21: E231 missing whitespace after ','
-./src/tfscreen/plot/plot_theta_fits.py:124:28: E231 missing whitespace after ','
-./src/tfscreen/plot/plot_theta_fits.py:125:1: W293 blank line contains whitespace
+./src/tfscreen/plot/plot_theta_fits.py:124:25: E231 missing whitespace after ','
./src/tfscreen/plot/plot_theta_fits.py:130:38: E231 missing whitespace after ','
-./src/tfscreen/plot/plot_theta_fits.py:133:27: E231 missing whitespace after ','
-./src/tfscreen/plot/plot_theta_fits.py:134:1: W293 blank line contains whitespace
+./src/tfscreen/plot/plot_theta_fits.py:133:25: E231 missing whitespace after ','
./src/tfscreen/plot/plot_theta_fits.py:139:38: E231 missing whitespace after ','
./src/tfscreen/plot/plot_theta_fits.py:141:1: W293 blank line contains whitespace
-./src/tfscreen/plot/plot_theta_fits.py:152:1: W391 blank line at end of file
+./src/tfscreen/plot/plot_theta_fits.py:144:22: E231 missing whitespace after ','
+./src/tfscreen/plot/plot_theta_fits.py:153:1: W391 blank line at end of file
./src/tfscreen/plot/uncertainty_calibration.py:12:1: E302 expected 2 blank lines, found 1
./src/tfscreen/plot/uncertainty_calibration.py:21:5: E125 continuation line with same indent as next logical line
./src/tfscreen/plot/uncertainty_calibration.py:24:63: W291 trailing whitespace
@@ -1268,23 +1264,54 @@
./src/tfscreen/process_raw/scripts/process_fastq_cli.py:418:33: E231 missing whitespace after ','
./src/tfscreen/process_raw/scripts/process_fastq_cli.py:420:1: E302 expected 2 blank lines, found 1
./src/tfscreen/simulate/__init__.py:29:2: W292 no newline at end of file
+./src/tfscreen/simulate/binding_params.py:225:1: C901 'build_theta_gc_override_hill_mut' is too complex (15)
+./src/tfscreen/simulate/binding_params.py:289:17: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:290:18: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:292:13: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:293:13: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:296:16: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:298:12: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:299:12: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:317:20: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:318:21: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:321:27: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:321:43: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:323:23: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:324:23: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:330:32: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:332:28: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:333:28: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:336:16: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:350:14: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:350:31: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:352:10: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:352:27: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:353:10: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:353:27: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:383:22: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:385:18: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:386:18: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:389:18: E221 multiple spaces before operator
+./src/tfscreen/simulate/binding_params.py:391:15: E221 multiple spaces before operator
./src/tfscreen/simulate/build_sample_dataframes.py:4:1: E302 expected 2 blank lines, found 1
./src/tfscreen/simulate/build_sample_dataframes.py:56:1: W293 blank line contains whitespace
./src/tfscreen/simulate/build_sample_dataframes.py:68:1: W293 blank line contains whitespace
./src/tfscreen/simulate/build_sample_dataframes.py:83:1: W293 blank line contains whitespace
./src/tfscreen/simulate/build_sample_dataframes.py:87:1: W293 blank line contains whitespace
./src/tfscreen/simulate/build_sample_dataframes.py:95:21: W292 no newline at end of file
-./src/tfscreen/simulate/library_prediction.py:16:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/simulate/library_prediction.py:17:43: E252 missing whitespace around parameter equals
-./src/tfscreen/simulate/library_prediction.py:17:44: E252 missing whitespace around parameter equals
+./src/tfscreen/simulate/library_prediction.py:28:1: C901 'library_prediction' is too complex (14)
+./src/tfscreen/simulate/library_prediction.py:28:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/simulate/library_prediction.py:29:43: E252 missing whitespace around parameter equals
+./src/tfscreen/simulate/library_prediction.py:29:44: E252 missing whitespace around parameter equals
+./src/tfscreen/simulate/library_prediction.py:123:9: F841 local variable 'binding_noise' is assigned to but never used
./src/tfscreen/simulate/run_simulation.py:12:1: E302 expected 2 blank lines, found 1
./src/tfscreen/simulate/run_simulation.py:17:1: W293 blank line contains whitespace
./src/tfscreen/simulate/run_simulation.py:18:36: E231 missing whitespace after ','
./src/tfscreen/simulate/run_simulation.py:21:1: W293 blank line contains whitespace
./src/tfscreen/simulate/run_simulation.py:22:23: E231 missing whitespace after ','
-./src/tfscreen/simulate/run_simulation.py:22:35: E231 missing whitespace after ','
-./src/tfscreen/simulate/run_simulation.py:22:52: E231 missing whitespace after ','
-./src/tfscreen/simulate/run_simulation.py:22:61: E231 missing whitespace after ','
+./src/tfscreen/simulate/run_simulation.py:22:36: E231 missing whitespace after ','
+./src/tfscreen/simulate/run_simulation.py:22:53: E231 missing whitespace after ','
+./src/tfscreen/simulate/run_simulation.py:22:62: E231 missing whitespace after ','
./src/tfscreen/simulate/run_simulation.py:24:47: E231 missing whitespace after ','
./src/tfscreen/simulate/run_simulation.py:25:32: E231 missing whitespace after ','
./src/tfscreen/simulate/run_simulation.py:34:1: W293 blank line contains whitespace
@@ -1306,7 +1333,7 @@
./src/tfscreen/simulate/run_simulation.py:88:51: E231 missing whitespace after ','
./src/tfscreen/simulate/run_simulation.py:88:62: E231 missing whitespace after ','
./src/tfscreen/simulate/run_simulation.py:91:26: E231 missing whitespace after ':'
-./src/tfscreen/simulate/run_simulation.py:92:28: E231 missing whitespace after ':'
+./src/tfscreen/simulate/run_simulation.py:92:29: E231 missing whitespace after ':'
./src/tfscreen/simulate/run_simulation.py:93:33: E231 missing whitespace after ':'
./src/tfscreen/simulate/run_simulation.py:94:25: E231 missing whitespace after ':'
./src/tfscreen/simulate/run_simulation.py:95:25: E231 missing whitespace after ':'
@@ -1315,155 +1342,177 @@
./src/tfscreen/simulate/run_simulation.py:101:1: W293 blank line contains whitespace
./src/tfscreen/simulate/run_simulation.py:106:1: W293 blank line contains whitespace
./src/tfscreen/simulate/run_simulation.py:107:1: W391 blank line at end of file
+./src/tfscreen/simulate/sample_theta.py:170:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:171:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:172:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:173:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:174:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:175:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:176:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:177:30: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/sample_theta.py:270:19: E221 multiple spaces before operator
./src/tfscreen/simulate/scripts/setup_sim_grid_cli.py:101:1: C901 'setup_sim_grid' is too complex (17)
-./src/tfscreen/simulate/selection_experiment.py:41:38: E231 missing whitespace after ':'
-./src/tfscreen/simulate/selection_experiment.py:42:37: E231 missing whitespace after ':'
-./src/tfscreen/simulate/selection_experiment.py:43:36: E231 missing whitespace after ':'
+./src/tfscreen/simulate/scripts/simulate_cli.py:119:15: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:120:17: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:121:16: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:122:16: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:126:20: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:127:22: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:128:21: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:147:43: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/scripts/simulate_cli.py:148:14: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:149:16: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:154:17: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:167:19: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:170:19: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:171:17: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:172:16: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:175:17: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:177:25: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:179:19: E221 multiple spaces before operator
+./src/tfscreen/simulate/scripts/simulate_cli.py:193:19: E221 multiple spaces before operator
+./src/tfscreen/simulate/selection_experiment.py:42:38: E231 missing whitespace after ':'
+./src/tfscreen/simulate/selection_experiment.py:43:37: E231 missing whitespace after ':'
./src/tfscreen/simulate/selection_experiment.py:44:36: E231 missing whitespace after ':'
./src/tfscreen/simulate/selection_experiment.py:45:36: E231 missing whitespace after ':'
-./src/tfscreen/simulate/selection_experiment.py:46:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:47:1: C901 '_check_dict_number' is too complex (13)
-./src/tfscreen/simulate/selection_experiment.py:47:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/simulate/selection_experiment.py:90:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:133:1: C901 '_check_cf' is too complex (26)
-./src/tfscreen/simulate/selection_experiment.py:155:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:175:57: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:177:48: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:177:51: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:177:65: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:177:85: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:200:28: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:201:28: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:267:1: C901 '_check_lib_spec' is too complex (12)
-./src/tfscreen/simulate/selection_experiment.py:267:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/simulate/selection_experiment.py:319:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:324:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:346:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:399:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:415:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:423:36: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:423:44: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:433:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/simulate/selection_experiment.py:443:78: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:486:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:490:77: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:491:75: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:501:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:504:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:550:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:561:38: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:562:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:565:53: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:569:55: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:575:32: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:580:30: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:581:53: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:585:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/simulate/selection_experiment.py:627:51: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:633:35: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:638:78: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:655:50: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:655:52: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:657:75: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:659:44: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:660:46: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:661:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:662:67: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:712:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:728:79: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:734:29: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:776:43: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:782:74: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:783:18: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:789:79: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:791:33: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:796:80: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:801:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:807:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:818:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:861:73: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:866:76: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:883:25: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:984:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:992:26: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:992:36: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:992:44: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:992:52: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:992:60: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1000:28: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1001:32: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1009:42: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1018:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1019:78: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1021:63: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1027:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1029:51: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1036:28: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1038:77: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1039:24: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1046:79: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1051:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1052:77: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1054:36: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1056:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1057:61: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1059:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1065:128: E501 line too long (145 > 127 characters)
-./src/tfscreen/simulate/selection_experiment.py:1066:128: E501 line too long (145 > 127 characters)
-./src/tfscreen/simulate/selection_experiment.py:1070:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1071:72: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1075:34: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1082:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1084:20: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1084:53: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1085:20: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1085:56: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1087:32: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1088:38: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1100:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1108:41: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1115:49: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1117:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1122:36: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1123:37: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1123:45: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1123:53: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1123:61: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1128:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/simulate/selection_experiment.py:1163:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1184:57: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1185:53: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1188:28: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1188:39: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1193:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1199:78: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1202:19: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1210:33: E261 at least two spaces before inline comment
-./src/tfscreen/simulate/selection_experiment.py:1211:49: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1212:37: E231 missing whitespace after ':'
-./src/tfscreen/simulate/selection_experiment.py:1214:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1215:76: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1217:19: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1219:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1227:77: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1228:43: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1232:41: E127 continuation line over-indented for visual indent
-./src/tfscreen/simulate/selection_experiment.py:1235:49: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1259:71: W291 trailing whitespace
-./src/tfscreen/simulate/selection_experiment.py:1267:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1268:76: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1273:11: F541 f-string is missing placeholders
-./src/tfscreen/simulate/selection_experiment.py:1278:38: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1295:6: E114 indentation is not a multiple of 4 (comment)
-./src/tfscreen/simulate/selection_experiment.py:1302:51: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1303:48: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1304:1: W293 blank line contains whitespace
-./src/tfscreen/simulate/selection_experiment.py:1305:33: E231 missing whitespace after ','
-./src/tfscreen/simulate/selection_experiment.py:1308:1: W391 blank line at end of file
+./src/tfscreen/simulate/selection_experiment.py:46:36: E231 missing whitespace after ':'
+./src/tfscreen/simulate/selection_experiment.py:71:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:72:1: C901 '_check_dict_number' is too complex (13)
+./src/tfscreen/simulate/selection_experiment.py:72:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/simulate/selection_experiment.py:115:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:158:1: C901 '_check_cf' is too complex (25)
+./src/tfscreen/simulate/selection_experiment.py:180:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:199:57: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:223:28: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:224:28: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:290:1: C901 '_check_lib_spec' is too complex (12)
+./src/tfscreen/simulate/selection_experiment.py:290:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/simulate/selection_experiment.py:342:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:347:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:369:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:422:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:438:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:446:36: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:446:44: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:456:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/simulate/selection_experiment.py:466:78: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:509:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:513:77: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:514:75: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:524:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:527:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:573:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:584:38: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:585:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:588:53: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:592:55: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:598:32: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:603:30: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:604:53: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:608:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/simulate/selection_experiment.py:650:51: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:656:35: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:661:78: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:678:50: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:678:52: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:680:75: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:682:44: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:683:46: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:684:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:685:67: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:735:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:752:79: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:758:29: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:800:43: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:806:74: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:807:18: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:813:79: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:815:33: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:820:80: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:825:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:831:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:842:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:885:73: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:890:76: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:907:25: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1015:26: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1015:36: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1015:44: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1015:52: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1015:60: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1023:28: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1024:32: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1032:42: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1041:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1042:78: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1044:63: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1050:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1052:51: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1059:28: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1061:77: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1062:24: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1069:79: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1074:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1075:77: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1077:36: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1079:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1083:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1089:128: E501 line too long (145 > 127 characters)
+./src/tfscreen/simulate/selection_experiment.py:1090:128: E501 line too long (145 > 127 characters)
+./src/tfscreen/simulate/selection_experiment.py:1094:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1095:72: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1099:34: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1106:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1108:20: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1108:53: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1109:20: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1111:32: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1112:38: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1124:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1132:41: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1139:49: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1141:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1146:36: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1147:37: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1147:45: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1147:53: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1147:61: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1152:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/simulate/selection_experiment.py:1187:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1208:57: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1209:53: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1212:28: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1212:39: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1217:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1223:78: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1226:19: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1234:33: E261 at least two spaces before inline comment
+./src/tfscreen/simulate/selection_experiment.py:1235:49: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1236:37: E231 missing whitespace after ':'
+./src/tfscreen/simulate/selection_experiment.py:1238:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1239:76: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1241:19: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1243:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1251:77: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1252:43: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1256:41: E127 continuation line over-indented for visual indent
+./src/tfscreen/simulate/selection_experiment.py:1259:49: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1283:71: W291 trailing whitespace
+./src/tfscreen/simulate/selection_experiment.py:1291:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1292:76: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1297:11: F541 f-string is missing placeholders
+./src/tfscreen/simulate/selection_experiment.py:1302:38: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1319:6: E114 indentation is not a multiple of 4 (comment)
+./src/tfscreen/simulate/selection_experiment.py:1326:51: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1327:48: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1328:1: W293 blank line contains whitespace
+./src/tfscreen/simulate/selection_experiment.py:1329:33: E231 missing whitespace after ','
+./src/tfscreen/simulate/selection_experiment.py:1332:1: W391 blank line at end of file
./src/tfscreen/simulate/sim_data_class.py:97:1: E303 too many blank lines (3)
./src/tfscreen/simulate/thermo_to_growth.py:102:36: E127 continuation line over-indented for visual indent
./src/tfscreen/simulate/thermo_to_growth.py:103:36: E127 continuation line over-indented for visual indent
-./src/tfscreen/simulate/thermo_to_growth.py:395:5: F841 local variable 'n_geno' is assigned to but never used
+./src/tfscreen/simulate/thermo_to_growth.py:289:1: C901 'thermo_to_growth' is too complex (18)
+./src/tfscreen/simulate/thermo_to_growth.py:626:41: E127 continuation line over-indented for visual indent
./src/tfscreen/tfmodel/__init__.py:2:73: W291 trailing whitespace
./src/tfscreen/tfmodel/__init__.py:3:20: W291 trailing whitespace
./src/tfscreen/tfmodel/__init__.py:18:65: W291 trailing whitespace
@@ -1477,12 +1526,11 @@
./src/tfscreen/tfmodel/__init__.py:41:77: W291 trailing whitespace
./src/tfscreen/tfmodel/__init__.py:43:56: W291 trailing whitespace
./src/tfscreen/tfmodel/__init__.py:46:1: W391 blank line at end of file
+./src/tfscreen/tfmodel/analysis/error_calibration.py:474:32: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/analysis/error_calibration.py:511:1: C901 'summarize_sbc' is too complex (16)
./src/tfscreen/tfmodel/analysis/extraction.py:81:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/analysis/extraction.py:85:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/analysis/extraction.py:103:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/analysis/extraction.py:116:72: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/extraction.py:117:71: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/extraction.py:118:70: W291 trailing whitespace
./src/tfscreen/tfmodel/analysis/extraction.py:182:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/analysis/extraction.py:260:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/analysis/extraction.py:261:29: E128 continuation line under-indented for visual indent
@@ -1497,79 +1545,50 @@
./src/tfscreen/tfmodel/analysis/predict_unmeasured.py:45:12: E221 multiple spaces before operator
./src/tfscreen/tfmodel/analysis/predict_unmeasured.py:99:16: E221 multiple spaces before operator
./src/tfscreen/tfmodel/analysis/predict_unmeasured.py:120:13: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/analysis/prediction.py:11:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/analysis/prediction.py:20:76: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:21:37: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:28:75: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:29:47: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:31:70: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:34:68: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:40:73: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:77:71: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:98:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:105:1: C901 'predict' is too complex (25)
-./src/tfscreen/tfmodel/analysis/prediction.py:105:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/analysis/prediction.py:166:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:177:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/analysis/prediction.py:179:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:192:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:194:18: E111 indentation is not a multiple of 4
-./src/tfscreen/tfmodel/analysis/prediction.py:194:18: E117 over-indented
-./src/tfscreen/tfmodel/analysis/prediction.py:203:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:207:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:216:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:228:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:231:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:236:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:239:67: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:240:76: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:244:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:251:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:260:58: E741 ambiguous variable name 'l'
-./src/tfscreen/tfmodel/analysis/prediction.py:263:26: E111 indentation is not a multiple of 4
-./src/tfscreen/tfmodel/analysis/prediction.py:263:26: E117 over-indented
-./src/tfscreen/tfmodel/analysis/prediction.py:266:25: E122 continuation line missing indentation or outdented
-./src/tfscreen/tfmodel/analysis/prediction.py:272:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:273:46: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:274:62: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:276:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:278:40: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:279:42: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:280:47: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:282:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:285:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:302:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:303:71: W291 trailing whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:307:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:312:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/analysis/prediction.py:316:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:12:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/analysis/prediction.py:71:1: C901 'copy_orchestrator' is too complex (11)
+./src/tfscreen/tfmodel/analysis/prediction.py:196:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/analysis/prediction.py:242:1: C901 'predict' is too complex (22)
+./src/tfscreen/tfmodel/analysis/prediction.py:303:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:369:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:372:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:377:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:380:67: W291 trailing whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:381:76: W291 trailing whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:385:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:392:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/analysis/prediction.py:401:58: E741 ambiguous variable name 'l'
+./src/tfscreen/tfmodel/analysis/prediction.py:404:26: E111 indentation is not a multiple of 4
+./src/tfscreen/tfmodel/analysis/prediction.py:404:26: E117 over-indented
+./src/tfscreen/tfmodel/analysis/prediction.py:407:25: E122 continuation line missing indentation or outdented
+./src/tfscreen/tfmodel/analysis/prediction.py:413:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/analysis/prior_predictive.py:104:33: E203 whitespace before ':'
-./src/tfscreen/tfmodel/analysis/sbc.py:130:1: C901 'summarize_sbc' is too complex (17)
-./src/tfscreen/tfmodel/configuration_io.py:9:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/configuration_io.py:27:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/configuration_io.py:106:1: C901 'write_configuration' is too complex (30)
-./src/tfscreen/tfmodel/configuration_io.py:106:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/configuration_io.py:177:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:180:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:182:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:183:40: E713 test for membership should be 'not in'
-./src/tfscreen/tfmodel/configuration_io.py:189:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:190:31: E713 test for membership should be 'not in'
-./src/tfscreen/tfmodel/configuration_io.py:196:18: W291 trailing whitespace
-./src/tfscreen/tfmodel/configuration_io.py:212:33: E261 at least two spaces before inline comment
-./src/tfscreen/tfmodel/configuration_io.py:224:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:238:1: E305 expected 2 blank lines after class or function definition, found 1
-./src/tfscreen/tfmodel/configuration_io.py:238:1: E402 module level import not at top of file
-./src/tfscreen/tfmodel/configuration_io.py:240:1: C901 'read_configuration' is too complex (19)
-./src/tfscreen/tfmodel/configuration_io.py:240:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/configuration_io.py:273:128: E501 line too long (139 > 127 characters)
-./src/tfscreen/tfmodel/configuration_io.py:278:22: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/configuration_io.py:279:22: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/configuration_io.py:280:22: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/configuration_io.py:286:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:298:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/configuration_io.py:313:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:19:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/configuration_io.py:37:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/configuration_io.py:116:1: C901 'write_configuration' is too complex (31)
+./src/tfscreen/tfmodel/configuration_io.py:116:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/configuration_io.py:190:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:193:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:195:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:196:40: E713 test for membership should be 'not in'
+./src/tfscreen/tfmodel/configuration_io.py:202:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:203:31: E713 test for membership should be 'not in'
+./src/tfscreen/tfmodel/configuration_io.py:209:18: W291 trailing whitespace
+./src/tfscreen/tfmodel/configuration_io.py:225:33: E261 at least two spaces before inline comment
+./src/tfscreen/tfmodel/configuration_io.py:237:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:251:1: E305 expected 2 blank lines after class or function definition, found 1
+./src/tfscreen/tfmodel/configuration_io.py:251:1: E402 module level import not at top of file
+./src/tfscreen/tfmodel/configuration_io.py:253:1: C901 'read_configuration' is too complex (20)
+./src/tfscreen/tfmodel/configuration_io.py:253:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/configuration_io.py:289:128: E501 line too long (139 > 127 characters)
+./src/tfscreen/tfmodel/configuration_io.py:297:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/configuration_io.py:298:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/configuration_io.py:299:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/configuration_io.py:300:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/configuration_io.py:306:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/configuration_io.py:318:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:333:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/configuration_io.py:338:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/data_class.py:12:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/data_class.py:27:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/data_class.py:28:5: E303 too many blank lines (2)
@@ -1577,11 +1596,9 @@
./src/tfscreen/tfmodel/data_class.py:104:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/data_class.py:121:49: W291 trailing whitespace
./src/tfscreen/tfmodel/data_class.py:123:76: W291 trailing whitespace
-./src/tfscreen/tfmodel/data_class.py:160:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/data_class.py:166:48: W291 trailing whitespace
-./src/tfscreen/tfmodel/data_class.py:185:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/data_class.py:192:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/data_class.py:193:5: E266 too many leading '#' for block comment
+./src/tfscreen/tfmodel/data_class.py:214:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/data_class.py:221:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/data_class.py:222:5: E266 too many leading '#' for block comment
./src/tfscreen/tfmodel/generative/components/__init__.py:3:4: W292 no newline at end of file
./src/tfscreen/tfmodel/generative/components/activity/fixed.py:8:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/generative/components/activity/fixed.py:16:1: E302 expected 2 blank lines, found 1
@@ -1798,19 +1815,21 @@
./src/tfscreen/tfmodel/generative/components/dk_geno/hierarchical_geno.py:374:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/generative/components/dk_geno/hierarchical_geno.py:374:5: W292 no newline at end of file
./src/tfscreen/tfmodel/generative/components/growth/linear.py:28:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:179:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:193:16: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:199:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:201:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:202:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:203:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:204:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:205:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:206:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:207:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:208:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:209:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/growth/linear.py:232:12: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:92:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:237:40: E261 at least two spaces before inline comment
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:260:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:274:16: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:280:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:282:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:283:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:284:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:285:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:286:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:287:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:288:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:289:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:290:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/growth/linear.py:313:12: E221 multiple spaces before operator
./src/tfscreen/tfmodel/generative/components/growth/power.py:25:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/generative/components/growth/saturation.py:23:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/generative/components/growth_transition/baranyi.py:133:30: E261 at least two spaces before inline comment
@@ -1940,151 +1959,220 @@
./src/tfscreen/tfmodel/generative/components/theta/categorical_geno.py:553:20: E221 multiple spaces before operator
./src/tfscreen/tfmodel/generative/components/theta/categorical_geno.py:567:15: E221 multiple spaces before operator
./src/tfscreen/tfmodel/generative/components/theta/categorical_geno.py:573:18: W292 no newline at end of file
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:105:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:111:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:112:12: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:115:12: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:121:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:123:8: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:124:8: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:127:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:131:7: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:204:29: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:206:25: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:207:25: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:212:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:212:53: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:212:74: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:214:21: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:215:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:216:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:219:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:221:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:257:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:262:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:266:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:271:22: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:275:16: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:280:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:284:16: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:289:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:294:28: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:298:30: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:302:24: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:306:24: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:314:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:317:22: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:320:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:323:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:329:29: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:333:25: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:335:25: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:342:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:342:53: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:342:74: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:344:21: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:345:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:346:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:348:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:350:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:390:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:393:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:396:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:397:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:423:48: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:424:50: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:425:46: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:427:50: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:428:52: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:429:48: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:431:49: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:432:51: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:433:47: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:435:49: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:436:51: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:437:47: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:464:13: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:472:43: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:473:45: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:474:45: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:476:44: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:478:44: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:481:40: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:483:41: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:484:41: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:569:24: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:571:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:572:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:575:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:577:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:578:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:590:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:618:40: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:680:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:683:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:135:76: E261 at least two spaces before inline comment
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:235:31: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:236:31: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:237:31: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:275:16: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:275:30: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:277:17: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:277:28: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:278:17: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:278:28: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:289:22: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:289:38: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:291:23: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:291:36: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:292:23: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:292:36: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:302:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:304:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:305:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:308:24: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:308:44: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:310:19: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:310:39: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:311:19: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:311:39: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:543:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:545:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:546:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:552:24: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:552:44: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:552:73: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:554:19: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:554:39: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:554:71: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:555:19: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:555:39: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:555:71: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:784:17: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:786:13: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:787:13: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:788:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:789:12: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:790:8: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:791:8: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:794:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:794:43: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:796:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:796:39: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:797:10: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:797:39: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:800:16: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:802:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:803:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:804:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:806:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:807:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:811:18: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:813:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:814:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:857:40: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:912:17: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:925:11: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:928:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:160:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:166:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:167:12: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:170:12: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:176:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:178:8: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:179:8: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:182:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:186:7: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:259:29: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:261:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:262:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:267:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:267:53: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:267:74: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:269:21: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:270:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:271:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:274:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:276:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:312:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:317:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:321:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:326:22: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:330:16: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:335:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:339:16: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:344:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:349:28: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:353:30: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:357:24: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:361:24: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:369:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:372:22: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:375:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:378:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:384:29: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:388:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:390:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:397:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:397:53: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:397:74: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:399:21: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:400:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:401:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:403:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:405:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:464:17: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:466:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:469:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:482:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:484:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:485:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:490:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:492:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:493:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:498:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:500:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:501:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:506:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:508:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:509:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:514:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:516:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:517:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:520:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:521:70: E261 at least two spaces before inline comment
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:522:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:526:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:529:18: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:529:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:529:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:530:19: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:530:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:531:19: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:531:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:532:15: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:532:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:532:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:533:11: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:533:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:533:16: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:534:14: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:534:16: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:565:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:568:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:571:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:572:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:598:48: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:599:50: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:600:46: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:602:50: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:603:52: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:604:48: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:606:49: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:607:51: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:608:47: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:610:49: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:611:51: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:612:47: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:632:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:666:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:674:43: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:675:45: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:676:45: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:678:44: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:680:44: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:683:40: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:685:41: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:686:41: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:771:24: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:773:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:774:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:777:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:779:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:780:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:792:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:820:40: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:882:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_geno.py:885:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:193:76: E261 at least two spaces before inline comment
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:300:31: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:301:31: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:302:31: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:340:16: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:340:30: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:342:17: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:342:28: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:343:17: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:343:28: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:354:22: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:354:38: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:356:23: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:356:36: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:357:23: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:357:36: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:367:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:369:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:370:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:373:24: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:373:44: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:375:19: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:375:39: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:376:19: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:376:39: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:608:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:610:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:611:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:617:24: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:617:44: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:617:73: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:619:19: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:619:39: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:619:71: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:620:19: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:620:39: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:620:71: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:700:17: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:702:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:705:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:712:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:712:31: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:714:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:715:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:715:27: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:748:22: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:750:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:751:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:755:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:756:70: E261 at least two spaces before inline comment
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:757:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:761:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:764:18: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:764:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:764:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:765:19: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:765:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:766:19: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:766:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:767:15: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:767:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:767:21: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:768:11: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:768:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:768:16: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:769:14: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:769:16: E251 unexpected spaces around keyword / parameter equals
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:846:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1013:17: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1015:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1016:13: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1017:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1018:12: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1019:8: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1020:8: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1023:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1023:43: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1025:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1025:39: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1026:10: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1026:39: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1029:16: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1031:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1032:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1033:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1035:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1036:14: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1040:18: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1042:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1043:15: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1086:40: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1141:17: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1154:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/components/theta/hill_mut.py:1157:14: E221 multiple spaces before operator
./src/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/PK.py:38:1: F401 'tfscreen.tfmodel.generative.components.theta.thermo.O2_C12_K5_U0_a.thermo.run_model' imported but unused
./src/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/PK.py:38:1: F401 'tfscreen.tfmodel.generative.components.theta.thermo.O2_C12_K5_U0_a.thermo.get_population_moments' imported but unused
./src/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/PK.py:171:6: E221 multiple spaces before operator
@@ -3148,7 +3236,6 @@
./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:462:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:463:54: W291 trailing whitespace
./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:464:79: W291 trailing whitespace
-./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:465:65: W291 trailing whitespace
./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:466:74: W291 trailing whitespace
./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:470:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/generative/components/transformation/_congression.py:473:1: E302 expected 2 blank lines, found 1
@@ -3173,18 +3260,20 @@
./src/tfscreen/tfmodel/generative/components/transformation/single.py:40:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/generative/components/transformation/single.py:46:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/generative/components/transformation/single.py:52:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/generative/model.py:11:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/generative/model.py:43:54: W291 trailing whitespace
-./src/tfscreen/tfmodel/generative/model.py:45:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/generative/model.py:103:37: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/generative/model.py:104:24: F541 f-string is missing placeholders
-./src/tfscreen/tfmodel/generative/model.py:104:45: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/generative/model.py:133:36: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/generative/model.py:134:24: F541 f-string is missing placeholders
-./src/tfscreen/tfmodel/generative/model.py:134:44: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/generative/model.py:213:28: F541 f-string is missing placeholders
-./src/tfscreen/tfmodel/generative/model.py:214:28: F541 f-string is missing placeholders
-./src/tfscreen/tfmodel/generative/model.py:220:1: W391 blank line at end of file
+./src/tfscreen/tfmodel/generative/model.py:12:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/generative/model.py:44:54: W291 trailing whitespace
+./src/tfscreen/tfmodel/generative/model.py:46:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/generative/model.py:104:37: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/generative/model.py:105:24: F541 f-string is missing placeholders
+./src/tfscreen/tfmodel/generative/model.py:105:45: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/generative/model.py:134:36: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/generative/model.py:135:24: F541 f-string is missing placeholders
+./src/tfscreen/tfmodel/generative/model.py:135:44: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/generative/model.py:195:19: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/model.py:196:19: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/generative/model.py:239:28: F541 f-string is missing placeholders
+./src/tfscreen/tfmodel/generative/model.py:240:28: F541 f-string is missing placeholders
+./src/tfscreen/tfmodel/generative/model.py:246:1: W391 blank line at end of file
./src/tfscreen/tfmodel/generative/observe/binding.py:9:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/generative/observe/binding.py:42:79: E231 missing whitespace after ','
./src/tfscreen/tfmodel/generative/observe/binding.py:43:83: E231 missing whitespace after ','
@@ -3203,69 +3292,68 @@
./src/tfscreen/tfmodel/generative/observe/growth.py:87:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/generative/observe/growth.py:95:5: F841 local variable 'nu' is assigned to but never used
./src/tfscreen/tfmodel/generative/observe/growth.py:97:11: W292 no newline at end of file
-./src/tfscreen/tfmodel/generative/registry.py:56:29: W291 trailing whitespace
-./src/tfscreen/tfmodel/generative/registry.py:57:28: W291 trailing whitespace
-./src/tfscreen/tfmodel/generative/registry.py:60:23: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:61:17: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:62:16: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:63:21: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:65:14: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:66:23: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:67:32: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:69:14: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:70:16: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:71:28: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:73:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:74:16: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:75:28: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:76:25: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:77:27: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:78:24: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:80:21: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:85:20: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:89:12: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:90:27: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:91:20: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:92:19: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:93:34: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:94:36: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:95:36: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:96:34: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:57:29: W291 trailing whitespace
+./src/tfscreen/tfmodel/generative/registry.py:58:28: W291 trailing whitespace
+./src/tfscreen/tfmodel/generative/registry.py:61:23: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:62:17: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:63:16: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:64:21: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:66:14: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:67:23: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:68:32: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:70:14: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:71:16: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:72:28: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:74:15: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:75:16: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:76:28: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:77:25: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:78:27: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:79:24: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:81:21: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:86:20: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:90:12: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:91:18: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:92:27: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:93:20: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:94:19: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:95:34: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:96:36: E231 missing whitespace after ':'
./src/tfscreen/tfmodel/generative/registry.py:97:36: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:98:36: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:99:35: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:100:37: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:101:37: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:102:35: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:98:34: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:99:36: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:100:36: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:101:35: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:102:37: E231 missing whitespace after ':'
./src/tfscreen/tfmodel/generative/registry.py:103:37: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:104:37: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:106:25: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:107:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:108:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:109:23: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:111:26: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:112:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:113:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:115:19: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:116:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:117:20: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:104:35: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:105:37: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:106:37: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:108:25: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:109:15: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:110:15: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:111:23: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:113:26: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:114:15: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:115:15: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:117:19: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:118:15: E231 missing whitespace after ':'
./src/tfscreen/tfmodel/generative/registry.py:119:20: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:120:15: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:121:17: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:123:24: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:124:18: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:125:17: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:121:20: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:122:15: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:123:17: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:125:24: E231 missing whitespace after ':'
./src/tfscreen/tfmodel/generative/registry.py:126:18: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:127:20: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:128:22: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:129:18: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:131:22: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/generative/registry.py:132:21: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/inference/posteriors.py:4:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/inference/posteriors.py:42:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/posteriors.py:50:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/inference/posteriors.py:84:17: F841 local variable 'e' is assigned to but never used
-./src/tfscreen/tfmodel/inference/posteriors.py:115:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/generative/registry.py:127:17: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:128:18: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:129:20: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:130:22: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:131:18: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:133:22: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/generative/registry.py:134:21: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/inference/posteriors.py:72:1: C901 'load_posteriors' is too complex (12)
+./src/tfscreen/tfmodel/inference/posteriors.py:72:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/inference/posteriors.py:117:17: F841 local variable 'e' is assigned to but never used
./src/tfscreen/tfmodel/inference/run_inference.py:21:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/inference/run_inference.py:28:1: E305 expected 2 blank lines after class or function definition, found 1
./src/tfscreen/tfmodel/inference/run_inference.py:28:1: E402 module level import not at top of file
@@ -3289,7 +3377,7 @@
./src/tfscreen/tfmodel/inference/run_inference.py:123:58: W291 trailing whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:124:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:129:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:132:5: C901 'RunInference.run_optimization' is too complex (17)
+./src/tfscreen/tfmodel/inference/run_inference.py:132:5: C901 'RunInference.run_optimization' is too complex (18)
./src/tfscreen/tfmodel/inference/run_inference.py:223:77: W291 trailing whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:224:73: W291 trailing whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:238:1: W293 blank line contains whitespace
@@ -3310,259 +3398,257 @@
./src/tfscreen/tfmodel/inference/run_inference.py:346:36: W291 trailing whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:347:50: W291 trailing whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:348:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:362:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:365:5: C901 'RunInference._get_genotype_dim_map' is too complex (13)
-./src/tfscreen/tfmodel/inference/run_inference.py:391:26: E261 at least two spaces before inline comment
-./src/tfscreen/tfmodel/inference/run_inference.py:430:5: C901 'RunInference.get_posteriors' is too complex (18)
-./src/tfscreen/tfmodel/inference/run_inference.py:473:59: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:489:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:370:5: C901 'RunInference._get_genotype_dim_map' is too complex (14)
+./src/tfscreen/tfmodel/inference/run_inference.py:396:26: E261 at least two spaces before inline comment
+./src/tfscreen/tfmodel/inference/run_inference.py:445:5: C901 'RunInference.get_posteriors' is too complex (18)
+./src/tfscreen/tfmodel/inference/run_inference.py:488:59: W291 trailing whitespace
./src/tfscreen/tfmodel/inference/run_inference.py:504:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:518:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:520:71: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:564:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:582:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:598:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:600:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/inference/run_inference.py:600:37: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:600:49: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:633:31: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:645:49: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:646:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:653:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/inference/run_inference.py:653:31: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:653:41: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:667:31: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/inference/run_inference.py:668:32: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/inference/run_inference.py:669:35: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/inference/run_inference.py:670:33: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/inference/run_inference.py:671:32: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/inference/run_inference.py:678:38: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:679:31: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:721:33: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:735:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:738:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:739:49: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:748:36: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:755:27: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:755:34: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:757:61: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:806:5: E303 too many blank lines (4)
-./src/tfscreen/tfmodel/inference/run_inference.py:847:26: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:847:33: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:850:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:862:41: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:864:32: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:864:42: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:915:5: C901 'RunInference.get_map_posteriors' is too complex (13)
-./src/tfscreen/tfmodel/inference/run_inference.py:1153:50: E127 continuation line over-indented for visual indent
-./src/tfscreen/tfmodel/inference/run_inference.py:1194:29: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:1197:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:1202:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:1210:35: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/inference/run_inference.py:1214:30: W291 trailing whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:1217:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:1221:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:1227:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/inference/run_inference.py:1230:5: C901 'RunInference.get_laplace_posteriors' is too complex (16)
-./src/tfscreen/tfmodel/inference/run_inference.py:1440:5: C901 'RunInference.get_nuts_posteriors' is too complex (12)
-./src/tfscreen/tfmodel/model_orchestrator.py:36:1: E303 too many blank lines (3)
-./src/tfscreen/tfmodel/model_orchestrator.py:70:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:73:39: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:77:25: E222 multiple spaces after operator
-./src/tfscreen/tfmodel/model_orchestrator.py:77:43: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:78:42: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:82:65: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:82:88: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:87:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:91:30: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:91:43: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:91:55: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:91:63: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:92:52: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:102:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:106:25: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:110:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:111:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/model_orchestrator.py:113:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/model_orchestrator.py:148:58: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:151:62: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:151:73: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:152:62: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:153:62: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:154:61: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:155:61: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:156:57: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:157:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:159:39: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:160:43: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:161:38: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:162:38: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:165:63: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:166:73: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:167:73: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:169:59: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:170:70: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:171:72: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:175:63: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:178:74: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:179:15: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:180:42: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:180:58: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:181:29: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/model_orchestrator.py:182:40: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:183:45: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:183:60: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:183:72: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:224:39: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:233:33: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:233:45: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:234:53: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:237:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/model_orchestrator.py:243:27: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:270:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/model_orchestrator.py:302:62: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:303:62: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:304:58: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:306:23: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:307:43: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:308:43: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:312:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:347:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:350:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:355:52: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:356:57: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:361:80: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:363:30: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:367:38: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:383:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:384:1: E303 too many blank lines (3)
-./src/tfscreen/tfmodel/model_orchestrator.py:404:82: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:409:87: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:412:84: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:417:79: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:497:5: C901 'ModelOrchestrator._initialize_data' is too complex (16)
-./src/tfscreen/tfmodel/model_orchestrator.py:497:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/model_orchestrator.py:522:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:523:28: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:525:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:526:35: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:526:48: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:526:56: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:527:46: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:527:66: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:528:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:532:39: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:533:34: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:538:67: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:545:32: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:547:10: E114 indentation is not a multiple of 4 (comment)
-./src/tfscreen/tfmodel/model_orchestrator.py:547:10: E116 unexpected indentation (comment)
-./src/tfscreen/tfmodel/model_orchestrator.py:549:76: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:552:45: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:554:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:556:49: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:568:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:576:57: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:598:38: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:599:41: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:599:56: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:600:44: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:600:60: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:601:40: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:601:58: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:602:46: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:602:71: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:604:77: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:611:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:616:39: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:616:45: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:616:53: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:623:40: E127 continuation line over-indented for visual indent
-./src/tfscreen/tfmodel/model_orchestrator.py:669:25: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/model_orchestrator.py:670:32: E127 continuation line over-indented for visual indent
-./src/tfscreen/tfmodel/model_orchestrator.py:671:26: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/model_orchestrator.py:672:32: E127 continuation line over-indented for visual indent
-./src/tfscreen/tfmodel/model_orchestrator.py:712:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:715:78: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:716:32: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:717:60: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:721:36: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:722:36: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:723:32: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:724:38: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:726:77: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:737:56: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:737:62: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:740:74: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:741:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:746:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:747:77: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:748:58: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:750:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:755:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:759:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:763:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:766:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:772:86: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:774:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:781:77: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:808:33: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:809:34: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:810:39: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:814:64: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:817:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:818:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:819:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/model_orchestrator.py:895:25: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/model_orchestrator.py:896:31: E127 continuation line over-indented for visual indent
-./src/tfscreen/tfmodel/model_orchestrator.py:898:31: E127 continuation line over-indented for visual indent
-./src/tfscreen/tfmodel/model_orchestrator.py:971:5: C901 'ModelOrchestrator._initialize_classes' is too complex (15)
-./src/tfscreen/tfmodel/model_orchestrator.py:991:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:993:67: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:994:71: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:996:68: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:997:77: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:998:75: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:999:40: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1016:42: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1017:43: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1018:40: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1018:43: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:1018:53: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1018:56: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/model_orchestrator.py:1018:64: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1066:75: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1069:69: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1073:75: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1075:69: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1085:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1087:9: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/model_orchestrator.py:1140:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1145:5: E303 too many blank lines (2)
-./src/tfscreen/tfmodel/model_orchestrator.py:1159:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1169:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1179:75: W291 trailing whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1193:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1196:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1225:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1237:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1249:1: W293 blank line contains whitespace
-./src/tfscreen/tfmodel/model_orchestrator.py:1255:25: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1256:27: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1257:31: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1258:32: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1259:22: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1260:22: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1261:23: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1262:20: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1263:29: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1264:28: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1265:33: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1266:34: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1267:27: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1268:28: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/model_orchestrator.py:1269:31: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:191:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/inference/run_inference.py:519:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:533:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:535:71: W291 trailing whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:579:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:597:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:613:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:615:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/inference/run_inference.py:615:37: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:615:49: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:648:31: W291 trailing whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:660:49: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:661:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:668:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/inference/run_inference.py:668:31: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:668:41: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:682:31: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/inference/run_inference.py:683:32: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/inference/run_inference.py:684:35: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/inference/run_inference.py:685:33: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/inference/run_inference.py:686:32: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/inference/run_inference.py:693:38: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:694:31: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:776:33: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:790:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:793:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:794:49: W291 trailing whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:803:36: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:810:27: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:810:34: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:812:61: W291 trailing whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:861:5: E303 too many blank lines (4)
+./src/tfscreen/tfmodel/inference/run_inference.py:902:26: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:902:33: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:905:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:917:41: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:919:32: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:919:42: W291 trailing whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:970:5: C901 'RunInference.get_map_posteriors' is too complex (13)
+./src/tfscreen/tfmodel/inference/run_inference.py:1208:50: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/inference/run_inference.py:1249:29: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:1252:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:1257:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:1265:35: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/inference/run_inference.py:1269:30: W291 trailing whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:1272:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:1276:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:1282:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/inference/run_inference.py:1285:5: C901 'RunInference.get_laplace_posteriors' is too complex (16)
+./src/tfscreen/tfmodel/inference/run_inference.py:1495:5: C901 'RunInference.get_nuts_posteriors' is too complex (12)
+./src/tfscreen/tfmodel/model_orchestrator.py:37:1: E303 too many blank lines (3)
+./src/tfscreen/tfmodel/model_orchestrator.py:71:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:74:39: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:78:25: E222 multiple spaces after operator
+./src/tfscreen/tfmodel/model_orchestrator.py:78:43: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:79:42: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:83:65: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:83:88: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:94:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:98:30: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:98:43: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:98:55: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:98:63: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:99:52: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:109:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:113:25: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:117:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:118:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/model_orchestrator.py:120:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/model_orchestrator.py:155:58: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:158:62: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:158:73: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:159:62: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:160:62: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:161:61: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:162:61: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:163:57: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:164:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:166:39: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:167:43: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:168:38: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:169:38: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:172:63: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:173:73: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:174:73: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:176:59: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:177:70: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:178:72: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:182:63: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:185:74: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:186:15: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:187:42: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:187:58: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:188:29: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/model_orchestrator.py:189:40: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:190:45: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:190:60: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:190:72: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:231:39: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:240:33: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:240:45: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:241:53: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:244:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/model_orchestrator.py:250:27: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:287:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/model_orchestrator.py:319:62: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:320:62: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:321:58: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:323:23: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:324:43: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:325:43: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:329:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:409:12: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/model_orchestrator.py:469:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:472:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:477:52: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:478:57: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:485:30: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:489:38: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:505:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:506:1: E303 too many blank lines (3)
+./src/tfscreen/tfmodel/model_orchestrator.py:526:82: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:531:87: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:534:84: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:539:79: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:621:5: C901 'ModelOrchestrator._initialize_data' is too complex (17)
+./src/tfscreen/tfmodel/model_orchestrator.py:621:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/model_orchestrator.py:646:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:647:28: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:649:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:650:35: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:650:48: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:650:56: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:651:46: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:651:66: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:652:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:656:39: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:657:34: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:662:67: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:669:32: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:671:10: E114 indentation is not a multiple of 4 (comment)
+./src/tfscreen/tfmodel/model_orchestrator.py:671:10: E116 unexpected indentation (comment)
+./src/tfscreen/tfmodel/model_orchestrator.py:673:76: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:676:45: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:678:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:680:49: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:692:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:700:57: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:722:38: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:723:41: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:723:56: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:724:44: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:724:60: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:725:40: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:725:58: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:726:46: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:726:71: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:728:77: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:735:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:740:39: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:740:45: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:740:53: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:747:40: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/model_orchestrator.py:793:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/model_orchestrator.py:794:32: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/model_orchestrator.py:795:26: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/model_orchestrator.py:796:32: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/model_orchestrator.py:836:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:839:78: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:840:32: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:841:60: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:845:36: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:846:36: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:847:32: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:848:38: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:850:77: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:861:56: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:861:62: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:864:74: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:865:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:870:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:871:77: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:872:58: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:874:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:879:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:883:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:887:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:890:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:896:86: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:898:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:905:77: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:961:33: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:962:34: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:963:35: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:964:39: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:971:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:972:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:973:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/model_orchestrator.py:1049:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/model_orchestrator.py:1050:31: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/model_orchestrator.py:1052:31: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/model_orchestrator.py:1125:5: C901 'ModelOrchestrator._initialize_classes' is too complex (16)
+./src/tfscreen/tfmodel/model_orchestrator.py:1145:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1147:67: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1148:71: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1150:68: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1151:77: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1152:75: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1153:40: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1170:42: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1171:43: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1172:40: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1172:43: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:1172:53: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1172:56: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/model_orchestrator.py:1172:64: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1229:75: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1232:69: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1236:75: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1238:69: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1248:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1250:9: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/model_orchestrator.py:1303:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1308:5: E303 too many blank lines (2)
+./src/tfscreen/tfmodel/model_orchestrator.py:1322:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1332:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1342:75: W291 trailing whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1356:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1359:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1388:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1400:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1412:1: W293 blank line contains whitespace
+./src/tfscreen/tfmodel/model_orchestrator.py:1418:25: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1419:27: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1420:31: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1421:32: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1422:22: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1423:22: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1424:23: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1425:20: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1426:29: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1427:28: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1428:33: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1429:34: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1430:27: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1431:28: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/model_orchestrator.py:1432:31: E231 missing whitespace after ':'
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:192:22: E128 continuation line under-indented for visual indent
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:193:22: E128 continuation line under-indented for visual indent
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:194:22: E128 continuation line under-indented for visual indent
@@ -3581,64 +3667,61 @@
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:207:22: E128 continuation line under-indented for visual indent
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:208:22: E128 continuation line under-indented for visual indent
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:209:22: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:219:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:221:59: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:222:58: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:223:55: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:224:60: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:210:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:211:22: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:223:1: E302 expected 2 blank lines, found 1
./src/tfscreen/tfmodel/scripts/configure_model_cli.py:225:59: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:226:63: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:227:55: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/configure_model_cli.py:229:1: E305 expected 2 blank lines after class or function definition, found 1
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:96:6: E114 indentation is not a multiple of 4 (comment)
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:96:6: E116 unexpected indentation (comment)
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:119:27: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:132:35: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:134:47: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:138:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:265:35: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:267:47: E231 missing whitespace after ','
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:271:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:547:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:549:60: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:550:53: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:551:64: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:552:66: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:553:66: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:554:64: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:555:65: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:556:64: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:557:72: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:558:74: E231 missing whitespace after ':'
-./src/tfscreen/tfmodel/scripts/fit_model_cli.py:560:1: E305 expected 2 blank lines after class or function definition, found 1
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:226:58: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:227:60: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:228:55: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:229:60: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:230:59: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:231:63: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:232:55: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/configure_model_cli.py:234:1: E305 expected 2 blank lines after class or function definition, found 1
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:97:6: E114 indentation is not a multiple of 4 (comment)
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:97:6: E116 unexpected indentation (comment)
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:120:27: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:124:35: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:126:47: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:130:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:245:39: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:247:51: E231 missing whitespace after ','
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:251:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:511:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:513:60: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:514:53: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:515:64: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:516:66: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:517:66: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:518:64: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:519:65: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:520:64: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:521:72: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:522:74: E231 missing whitespace after ':'
+./src/tfscreen/tfmodel/scripts/fit_model_cli.py:524:1: E305 expected 2 blank lines after class or function definition, found 1
./src/tfscreen/tfmodel/scripts/predict_theta_cli.py:12:1: C901 'predict_theta' is too complex (11)
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:254:24: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:255:24: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:256:24: E128 continuation line under-indented for visual indent
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:259:1: C901 '_inject_calibration_priors' is too complex (19)
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:441:1: C901 '_build_csv_updates' is too complex (13)
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:543:1: C901 '_apply_guesses_updates' is too complex (12)
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:698:1: C901 '_make_calibration_plots' is too complex (23)
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:784:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:787:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:788:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:839:7: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:840:7: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:843:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:844:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:845:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:846:14: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:889:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:890:15: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:912:19: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:927:21: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:950:20: E221 multiple spaces before operator
-./src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py:963:25: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:106:23: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:107:23: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:198:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:199:11: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:204:20: E221 multiple spaces before operator
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:268:24: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:269:24: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:270:24: E128 continuation line under-indented for visual indent
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:273:1: C901 '_inject_calibration_priors' is too complex (21)
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:464:1: C901 '_build_csv_updates' is too complex (11)
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:522:35: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:523:35: E127 continuation line over-indented for visual indent
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:662:1: C901 '_apply_guesses_updates' is too complex (12)
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:814:1: E303 too many blank lines (3)
+./src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py:818:1: E302 expected 2 blank lines, found 3
./src/tfscreen/tfmodel/scripts/setup_grid_cli.py:160:1: C901 'setup_grid' is too complex (17)
./src/tfscreen/tfmodel/scripts/subset_genotypes_cli.py:40:1: C901 'subset_genotypes' is too complex (23)
./src/tfscreen/tfmodel/scripts/subset_genotypes_cli.py:93:5: F841 local variable 'all_genotypes' is assigned to but never used
-./src/tfscreen/tfmodel/scripts/summarize_fit_cli.py:43:1: C901 '_read_all_losses' is too complex (11)
-./src/tfscreen/tfmodel/scripts/summarize_fit_cli.py:125:1: C901 'summarize_fit' is too complex (45)
+./src/tfscreen/tfmodel/scripts/summarize_fit_cli.py:154:1: C901 '_read_all_losses' is too complex (11)
+./src/tfscreen/tfmodel/scripts/summarize_fit_cli.py:413:1: C901 '_summarize_params' is too complex (23)
+./src/tfscreen/tfmodel/scripts/summarize_fit_cli.py:521:1: C901 'summarize_fit' is too complex (59)
./src/tfscreen/tfmodel/scripts/summarize_grid_cli.py:47:30: E127 continuation line over-indented for visual indent
./src/tfscreen/tfmodel/scripts/summarize_grid_cli.py:76:1: C901 'summarize_grid' is too complex (14)
./src/tfscreen/tfmodel/tensors/batch.py:5:1: E302 expected 2 blank lines, found 1
@@ -3723,7 +3806,7 @@
./src/tfscreen/tfmodel/tensors/tensor_manager.py:512:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/tensors/tensor_manager.py:517:1: W293 blank line contains whitespace
./src/tfscreen/tfmodel/tensors/tensor_manager.py:528:1: W391 blank line at end of file
-./src/tfscreen/util/__init__.py:70:2: W292 no newline at end of file
+./src/tfscreen/util/__init__.py:71:2: W292 no newline at end of file
./src/tfscreen/util/cli/__init__.py:16:1: W391 blank line at end of file
./src/tfscreen/util/cli/generalized_main.py:1:11: W291 trailing whitespace
./src/tfscreen/util/cli/generalized_main.py:5:1: C901 'generalized_main' is too complex (15)
@@ -3852,10 +3935,10 @@
./src/tfscreen/util/numerical/zero_truncated_poisson.py:71:34: W291 trailing whitespace
./src/tfscreen/util/numerical/zero_truncated_poisson.py:73:1: W293 blank line contains whitespace
./src/tfscreen/util/numerical/zero_truncated_poisson.py:74:23: W292 no newline at end of file
-./src/tfscreen/util/validation/check.py:7:1: C901 'check_number' is too complex (12)
-./src/tfscreen/util/validation/check.py:7:1: E302 expected 2 blank lines, found 1
-./src/tfscreen/util/validation/check.py:62:1: W293 blank line contains whitespace
-./src/tfscreen/util/validation/check.py:88:18: W292 no newline at end of file
+./src/tfscreen/util/validation/check.py:37:1: C901 'check_number' is too complex (12)
+./src/tfscreen/util/validation/check.py:37:1: E302 expected 2 blank lines, found 1
+./src/tfscreen/util/validation/check.py:92:1: W293 blank line contains whitespace
+./src/tfscreen/util/validation/check.py:118:18: W292 no newline at end of file
./tests/conftest.py:9:1: E302 expected 2 blank lines, found 1
./tests/conftest.py:40:63: E231 missing whitespace after ','
./tests/conftest.py:54:49: E231 missing whitespace after ','
@@ -3909,24 +3992,12 @@
./tests/smoke-tests/conftest.py:10:1: E302 expected 2 blank lines, found 1
./tests/smoke-tests/conftest.py:15:1: E302 expected 2 blank lines, found 1
./tests/smoke-tests/conftest.py:20:1: E302 expected 2 blank lines, found 1
-./tests/smoke-tests/test_configure_run_smoke.py:9:1: E302 expected 2 blank lines, found 1
-./tests/smoke-tests/test_configure_run_smoke.py:32:31: E127 continuation line over-indented for visual indent
+./tests/smoke-tests/test_configure_run_smoke.py:10:1: E302 expected 2 blank lines, found 1
./tests/smoke-tests/test_configure_run_smoke.py:33:31: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:51:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:52:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:53:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:54:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:55:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:56:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:97:31: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:98:31: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:122:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:123:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:124:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:125:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:126:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:127:25: E127 continuation line over-indented for visual indent
-./tests/smoke-tests/test_configure_run_smoke.py:170:16: E221 multiple spaces before operator
+./tests/smoke-tests/test_configure_run_smoke.py:34:31: E127 continuation line over-indented for visual indent
+./tests/smoke-tests/test_configure_run_smoke.py:101:31: E127 continuation line over-indented for visual indent
+./tests/smoke-tests/test_configure_run_smoke.py:102:31: E127 continuation line over-indented for visual indent
+./tests/smoke-tests/test_configure_run_smoke.py:175:16: E221 multiple spaces before operator
./tests/smoke-tests/test_model_smoke.py:3:1: F401 'shutil' imported but unused
./tests/smoke-tests/test_model_smoke.py:4:1: F401 'jax.numpy as jnp' imported but unused
./tests/smoke-tests/test_model_smoke.py:11:71: W291 trailing whitespace
@@ -4072,8 +4143,9 @@
./tests/smoke-tests/test_run_growth_analysis_smoke.py:37:1: W293 blank line contains whitespace
./tests/smoke-tests/test_run_growth_analysis_smoke.py:45:1: W293 blank line contains whitespace
./tests/smoke-tests/test_run_growth_analysis_smoke.py:48:1: W293 blank line contains whitespace
-./tests/smoke-tests/test_run_growth_analysis_smoke.py:62:1: W293 blank line contains whitespace
-./tests/smoke-tests/test_run_growth_analysis_smoke.py:66:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/cat_response/scripts/test_cat_response_cli.py:14:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/cat_response/scripts/test_cat_response_cli.py:19:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/cat_response/scripts/test_cat_response_cli.py:24:1: E305 expected 2 blank lines after class or function definition, found 1
./tests/tfscreen/analysis/cat_response/test_cat_fit.py:4:1: F401 'pandas as pd' imported but unused
./tests/tfscreen/analysis/cat_response/test_cat_fit.py:10:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/analysis/cat_response/test_cat_fit.py:17:1: E302 expected 2 blank lines, found 1
@@ -4172,12 +4244,6 @@
./tests/tfscreen/analysis/cat_response/test_cat_response.py:113:59: E261 at least two spaces before inline comment
./tests/tfscreen/analysis/cat_response/test_cat_response.py:114:1: W293 blank line contains whitespace
./tests/tfscreen/analysis/cat_response/test_cat_response.py:119:1: W391 blank line at end of file
-./tests/tfscreen/analysis/cat_response/test_cat_response_cli.py:14:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/analysis/cat_response/test_cat_response_cli.py:19:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/analysis/cat_response/test_cat_response_cli.py:24:1: E305 expected 2 blank lines after class or function definition, found 1
-./tests/tfscreen/analysis/cat_response/test_fit_response_cli.py:14:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/analysis/cat_response/test_fit_response_cli.py:19:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/analysis/cat_response/test_fit_response_cli.py:24:1: E305 expected 2 blank lines after class or function definition, found 1
./tests/tfscreen/analysis/test_extract_epistasis.py:16:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/analysis/test_extract_epistasis.py:22:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/analysis/test_extract_epistasis.py:34:35: W291 trailing whitespace
@@ -4229,6 +4295,33 @@
./tests/tfscreen/analysis/test_extract_epistasis.py:226:1: W293 blank line contains whitespace
./tests/tfscreen/analysis/test_extract_epistasis.py:228:74: E231 missing whitespace after ','
./tests/tfscreen/analysis/test_extract_epistasis.py:229:55: W292 no newline at end of file
+./tests/tfscreen/analysis/test_stats_test_suite.py:6:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/test_stats_test_suite.py:9:52: E261 at least two spaces before inline comment
+./tests/tfscreen/analysis/test_stats_test_suite.py:10:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:13:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:20:41: W291 trailing whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:23:39: W291 trailing whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:24:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:31:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:32:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/test_stats_test_suite.py:38:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:45:49: E261 at least two spaces before inline comment
+./tests/tfscreen/analysis/test_stats_test_suite.py:46:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:53:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/test_stats_test_suite.py:57:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:64:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/test_stats_test_suite.py:68:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:74:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/analysis/test_stats_test_suite.py:86:1: F401 'warnings' imported but unused
+./tests/tfscreen/analysis/test_stats_test_suite.py:86:1: E305 expected 2 blank lines after class or function definition, found 1
+./tests/tfscreen/analysis/test_stats_test_suite.py:86:1: E402 module level import not at top of file
+./tests/tfscreen/analysis/test_stats_test_suite.py:87:1: E302 expected 2 blank lines, found 0
+./tests/tfscreen/analysis/test_stats_test_suite.py:90:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:94:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:96:10: E111 indentation is not a multiple of 4
+./tests/tfscreen/analysis/test_stats_test_suite.py:96:10: E117 over-indented
+./tests/tfscreen/analysis/test_stats_test_suite.py:97:1: W293 blank line contains whitespace
+./tests/tfscreen/analysis/test_stats_test_suite.py:99:1: W391 blank line at end of file
./tests/tfscreen/genetics/test_build_cycles.py:7:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/genetics/test_build_cycles.py:21:38: E231 missing whitespace after ','
./tests/tfscreen/genetics/test_build_cycles.py:23:28: E124 closing bracket does not match visual indentation
@@ -4356,104 +4449,107 @@
./tests/tfscreen/genetics/test_library_manager.py:169:24: E231 missing whitespace after ':'
./tests/tfscreen/genetics/test_library_manager.py:173:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/genetics/test_library_manager.py:186:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:194:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:199:39: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:203:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:211:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:219:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:224:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:229:50: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:233:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:191:27: E221 multiple spaces before operator
+./tests/tfscreen/genetics/test_library_manager.py:192:32: E221 multiple spaces before operator
+./tests/tfscreen/genetics/test_library_manager.py:198:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:206:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:211:39: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:215:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:223:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:231:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:236:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:241:50: E261 at least two spaces before inline comment
./tests/tfscreen/genetics/test_library_manager.py:245:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:281:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:300:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:307:14: E201 whitespace after '('
-./tests/tfscreen/genetics/test_library_manager.py:307:39: E202 whitespace before ')'
-./tests/tfscreen/genetics/test_library_manager.py:308:14: E201 whitespace after '('
-./tests/tfscreen/genetics/test_library_manager.py:308:55: E202 whitespace before ')'
-./tests/tfscreen/genetics/test_library_manager.py:323:22: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:324:23: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:344:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:348:26: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:349:21: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:355:36: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:359:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:365:26: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:366:21: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:372:35: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:373:36: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:375:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:378:36: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:379:26: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:380:21: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:388:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:391:38: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:392:36: E114 indentation is not a multiple of 4 (comment)
-./tests/tfscreen/genetics/test_library_manager.py:392:36: E116 unexpected indentation (comment)
-./tests/tfscreen/genetics/test_library_manager.py:393:36: E114 indentation is not a multiple of 4 (comment)
-./tests/tfscreen/genetics/test_library_manager.py:393:36: E116 unexpected indentation (comment)
-./tests/tfscreen/genetics/test_library_manager.py:394:36: E114 indentation is not a multiple of 4 (comment)
-./tests/tfscreen/genetics/test_library_manager.py:394:36: E116 unexpected indentation (comment)
-./tests/tfscreen/genetics/test_library_manager.py:395:39: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:396:26: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:397:21: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:405:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:409:26: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:409:37: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:410:21: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:463:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:479:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:481:32: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:514:19: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:515:58: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:516:18: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:522:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:527:51: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:528:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:539:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:544:19: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:545:67: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:546:18: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:550:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:554:36: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:567:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:574:19: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:576:35: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:585:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:587:24: E201 whitespace after '{'
-./tests/tfscreen/genetics/test_library_manager.py:587:75: E202 whitespace before '}'
-./tests/tfscreen/genetics/test_library_manager.py:589:24: E201 whitespace after '{'
-./tests/tfscreen/genetics/test_library_manager.py:589:75: E202 whitespace before '}'
-./tests/tfscreen/genetics/test_library_manager.py:591:24: E201 whitespace after '{'
-./tests/tfscreen/genetics/test_library_manager.py:591:75: E202 whitespace before '}'
-./tests/tfscreen/genetics/test_library_manager.py:598:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:613:19: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:615:18: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:636:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:642:19: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:644:18: E222 multiple spaces after operator
-./tests/tfscreen/genetics/test_library_manager.py:644:44: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:653:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:668:1: E303 too many blank lines (3)
-./tests/tfscreen/genetics/test_library_manager.py:672:1: E302 expected 2 blank lines, found 3
-./tests/tfscreen/genetics/test_library_manager.py:681:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:689:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:691:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:694:48: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:695:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:707:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:711:42: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:712:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:715:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:719:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:724:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:728:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/genetics/test_library_manager.py:734:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:739:1: W293 blank line contains whitespace
-./tests/tfscreen/genetics/test_library_manager.py:771:22: E261 at least two spaces before inline comment
-./tests/tfscreen/genetics/test_library_manager.py:775:35: E231 missing whitespace after ':'
-./tests/tfscreen/genetics/test_library_manager.py:775:49: W291 trailing whitespace
-./tests/tfscreen/genetics/test_library_manager.py:800:34: W292 no newline at end of file
+./tests/tfscreen/genetics/test_library_manager.py:257:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:293:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:312:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:319:14: E201 whitespace after '('
+./tests/tfscreen/genetics/test_library_manager.py:319:39: E202 whitespace before ')'
+./tests/tfscreen/genetics/test_library_manager.py:320:14: E201 whitespace after '('
+./tests/tfscreen/genetics/test_library_manager.py:320:55: E202 whitespace before ')'
+./tests/tfscreen/genetics/test_library_manager.py:335:22: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:336:23: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:356:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:360:26: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:361:21: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:367:36: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:371:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:377:26: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:378:21: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:384:35: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:385:36: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:387:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:390:36: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:391:26: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:392:21: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:400:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:403:38: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:404:36: E114 indentation is not a multiple of 4 (comment)
+./tests/tfscreen/genetics/test_library_manager.py:404:36: E116 unexpected indentation (comment)
+./tests/tfscreen/genetics/test_library_manager.py:405:36: E114 indentation is not a multiple of 4 (comment)
+./tests/tfscreen/genetics/test_library_manager.py:405:36: E116 unexpected indentation (comment)
+./tests/tfscreen/genetics/test_library_manager.py:406:36: E114 indentation is not a multiple of 4 (comment)
+./tests/tfscreen/genetics/test_library_manager.py:406:36: E116 unexpected indentation (comment)
+./tests/tfscreen/genetics/test_library_manager.py:407:39: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:408:26: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:409:21: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:417:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:421:26: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:421:37: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:422:21: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:475:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:491:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:493:32: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:526:19: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:527:58: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:528:18: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:534:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:539:51: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:540:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:551:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:556:19: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:557:67: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:558:18: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:562:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:566:36: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:579:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:586:19: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:588:35: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:597:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:599:24: E201 whitespace after '{'
+./tests/tfscreen/genetics/test_library_manager.py:599:75: E202 whitespace before '}'
+./tests/tfscreen/genetics/test_library_manager.py:601:24: E201 whitespace after '{'
+./tests/tfscreen/genetics/test_library_manager.py:601:75: E202 whitespace before '}'
+./tests/tfscreen/genetics/test_library_manager.py:603:24: E201 whitespace after '{'
+./tests/tfscreen/genetics/test_library_manager.py:603:75: E202 whitespace before '}'
+./tests/tfscreen/genetics/test_library_manager.py:610:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:625:19: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:627:18: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:648:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:654:19: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:656:18: E222 multiple spaces after operator
+./tests/tfscreen/genetics/test_library_manager.py:656:44: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:665:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:680:1: E303 too many blank lines (3)
+./tests/tfscreen/genetics/test_library_manager.py:684:1: E302 expected 2 blank lines, found 3
+./tests/tfscreen/genetics/test_library_manager.py:694:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:702:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:704:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:707:48: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:708:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:720:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:734:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:747:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:751:42: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:759:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:764:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:768:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/genetics/test_library_manager.py:774:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:779:1: W293 blank line contains whitespace
+./tests/tfscreen/genetics/test_library_manager.py:811:22: E261 at least two spaces before inline comment
+./tests/tfscreen/genetics/test_library_manager.py:815:35: E231 missing whitespace after ':'
+./tests/tfscreen/genetics/test_library_manager.py:815:49: W291 trailing whitespace
+./tests/tfscreen/genetics/test_library_manager.py:840:34: W292 no newline at end of file
./tests/tfscreen/mle/curve_models/test_guesses.py:7:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/mle/curve_models/test_guesses.py:18:1: W293 blank line contains whitespace
./tests/tfscreen/mle/curve_models/test_guesses.py:21:44: E261 at least two spaces before inline comment
@@ -4678,33 +4774,6 @@
./tests/tfscreen/mle/test_predict_with_error.py:59:1: W293 blank line contains whitespace
./tests/tfscreen/mle/test_predict_with_error.py:64:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/mle/test_predict_with_error.py:67:61: W291 trailing whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:6:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/mle/test_stats_test_suite.py:9:52: E261 at least two spaces before inline comment
-./tests/tfscreen/mle/test_stats_test_suite.py:10:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:13:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:20:41: W291 trailing whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:23:39: W291 trailing whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:24:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:31:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:32:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/mle/test_stats_test_suite.py:38:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:45:49: E261 at least two spaces before inline comment
-./tests/tfscreen/mle/test_stats_test_suite.py:46:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:53:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/mle/test_stats_test_suite.py:57:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:64:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/mle/test_stats_test_suite.py:68:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:74:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/mle/test_stats_test_suite.py:86:1: F401 'warnings' imported but unused
-./tests/tfscreen/mle/test_stats_test_suite.py:86:1: E305 expected 2 blank lines after class or function definition, found 1
-./tests/tfscreen/mle/test_stats_test_suite.py:86:1: E402 module level import not at top of file
-./tests/tfscreen/mle/test_stats_test_suite.py:87:1: E302 expected 2 blank lines, found 0
-./tests/tfscreen/mle/test_stats_test_suite.py:90:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:94:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:96:10: E111 indentation is not a multiple of 4
-./tests/tfscreen/mle/test_stats_test_suite.py:96:10: E117 over-indented
-./tests/tfscreen/mle/test_stats_test_suite.py:97:1: W293 blank line contains whitespace
-./tests/tfscreen/mle/test_stats_test_suite.py:99:1: W391 blank line at end of file
./tests/tfscreen/plot/heatmap/test_heatmap_core.py:25:1: F401 'tfscreen.plot.default_styles.DEFAULT_HMAP_PATCH_KWARGS' imported but unused
./tests/tfscreen/plot/heatmap/test_heatmap_core.py:31:1: F401 'tfscreen.plot.helper.get_ax_limits' imported but unused
./tests/tfscreen/plot/heatmap/test_heatmap_core.py:31:47: W291 trailing whitespace
@@ -4880,6 +4949,16 @@
./tests/tfscreen/plot/test_est_v_real_summary.py:22:1: W293 blank line contains whitespace
./tests/tfscreen/plot/test_est_v_real_summary.py:25:1: W293 blank line contains whitespace
./tests/tfscreen/plot/test_est_v_real_summary.py:30:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/plot/test_geno_trajectory.py:18:1: E402 module level import not at top of file
+./tests/tfscreen/plot/test_geno_trajectory.py:20:1: E402 module level import not at top of file
+./tests/tfscreen/plot/test_geno_trajectory.py:159:25: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:304:31: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:305:42: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:320:31: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:321:42: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:488:19: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:516:19: E741 ambiguous variable name 'l'
+./tests/tfscreen/plot/test_geno_trajectory.py:544:19: E741 ambiguous variable name 'l'
./tests/tfscreen/plot/test_helper.py:3:1: F401 'pytest' imported but unused
./tests/tfscreen/plot/test_helper.py:6:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/plot/test_helper.py:9:1: W293 blank line contains whitespace
@@ -5023,73 +5102,73 @@
./tests/tfscreen/process_raw/test_fastq_to_counts.py:1091:1: W293 blank line contains whitespace
./tests/tfscreen/process_raw/test_fastq_to_counts.py:1099:1: W293 blank line contains whitespace
./tests/tfscreen/process_raw/test_fastq_to_counts.py:1104:45: W292 no newline at end of file
-./tests/tfscreen/process_raw/test_process_counts.py:16:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:25:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:44:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:68:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:82:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:89:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:96:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:102:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_counts.py:108:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:2:1: F401 'unittest.mock.call' imported but unused
-./tests/tfscreen/process_raw/test_process_fastq.py:22:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:31:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:38:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:39:26: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:33: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:40: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:47: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:54: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:61: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:68: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:75: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:82: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:39:89: E231 missing whitespace after ':'
-./tests/tfscreen/process_raw/test_process_fastq.py:44:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:55:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:72:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:77:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:84:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:99:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:110:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:117:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:122:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:135:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:149:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:151:62: E231 missing whitespace after ','
-./tests/tfscreen/process_raw/test_process_fastq.py:153:62: E231 missing whitespace after ','
-./tests/tfscreen/process_raw/test_process_fastq.py:154:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:177:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:183:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:193:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:200:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:209:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:215:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:218:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:233:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:240:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:250:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:253:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:257:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:260:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:267:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:273:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:280:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:284:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:288:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:298:38: E261 at least two spaces before inline comment
-./tests/tfscreen/process_raw/test_process_fastq.py:300:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:310:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:317:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:321:128: E501 line too long (132 > 127 characters)
-./tests/tfscreen/process_raw/test_process_fastq.py:332:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:343:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:348:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:350:26: E261 at least two spaces before inline comment
-./tests/tfscreen/process_raw/test_process_fastq.py:357:1: W293 blank line contains whitespace
-./tests/tfscreen/process_raw/test_process_fastq.py:360:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/process_raw/test_process_fastq.py:368:78: W292 no newline at end of file
+./tests/tfscreen/process_raw/test_process_counts_cli.py:16:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:25:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:44:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:68:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:82:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:89:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:96:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:102:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_counts_cli.py:108:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:2:1: F401 'unittest.mock.call' imported but unused
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:22:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:31:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:38:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:26: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:33: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:40: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:47: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:54: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:61: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:68: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:75: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:82: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:39:89: E231 missing whitespace after ':'
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:44:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:55:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:72:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:77:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:84:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:99:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:110:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:117:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:122:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:135:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:149:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:151:62: E231 missing whitespace after ','
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:153:62: E231 missing whitespace after ','
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:154:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:177:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:183:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:193:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:200:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:209:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:215:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:218:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:233:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:240:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:250:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:253:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:257:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:260:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:267:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:273:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:280:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:284:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:288:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:298:38: E261 at least two spaces before inline comment
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:300:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:310:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:317:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:321:128: E501 line too long (132 > 127 characters)
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:332:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:343:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:348:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:350:26: E261 at least two spaces before inline comment
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:357:1: W293 blank line contains whitespace
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:360:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/process_raw/test_process_fastq_cli.py:368:78: W292 no newline at end of file
./tests/tfscreen/simulate/growth/test_growth_linkage.py:50:16: E221 multiple spaces before operator
./tests/tfscreen/simulate/growth/test_growth_linkage.py:95:16: E221 multiple spaces before operator
./tests/tfscreen/simulate/growth/test_growth_linkage.py:154:16: E221 multiple spaces before operator
@@ -5121,6 +5200,37 @@
./tests/tfscreen/simulate/growth/test_transition_linkage.py:566:14: E221 multiple spaces before operator
./tests/tfscreen/simulate/growth/test_transition_linkage.py:567:14: E221 multiple spaces before operator
./tests/tfscreen/simulate/growth/test_transition_linkage.py:568:14: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:5:1: F401 'io' imported but unused
+./tests/tfscreen/simulate/test_binding_params.py:8:1: F401 'pandas as pd' imported but unused
+./tests/tfscreen/simulate/test_binding_params.py:11:1: F401 'tfscreen.simulate.binding_params._logit' imported but unused
+./tests/tfscreen/simulate/test_binding_params.py:83:20: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:85:16: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:86:17: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:87:25: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:89:21: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:90:21: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:91:21: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:93:19: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:203:22: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:205:32: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:215:35: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:217:32: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:229:35: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:237:35: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:302:38: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:304:37: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:313:5: F811 redefinition of unused 'pd' from line 8
+./tests/tfscreen/simulate/test_binding_params.py:316:15: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:317:14: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:422:43: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_binding_params.py:442:56: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_binding_params.py:443:56: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_binding_params.py:463:46: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_binding_params.py:484:38: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:485:39: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:487:37: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:505:45: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_binding_params.py:584:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_build_sample_dataframes.py:9:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_build_sample_dataframes.py:33:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_build_sample_dataframes.py:52:1: W293 blank line contains whitespace
@@ -5134,39 +5244,63 @@
./tests/tfscreen/simulate/test_build_sample_dataframes.py:104:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_build_sample_dataframes.py:115:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_build_sample_dataframes.py:128:55: W292 no newline at end of file
+./tests/tfscreen/simulate/test_generate_presplit_data.py:11:1: F401 'pytest' imported but unused
+./tests/tfscreen/simulate/test_generate_presplit_data.py:14:1: F401 'unittest.mock.MagicMock' imported but unused
+./tests/tfscreen/simulate/test_generate_presplit_data.py:74:39: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:86:42: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:87:42: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:89:39: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:97:39: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:105:39: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:123:39: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:124:14: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_generate_presplit_data.py:140:20: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_generate_presplit_data.py:141:51: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:142:22: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_generate_presplit_data.py:143:51: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:145:22: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_generate_presplit_data.py:158:35: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:160:35: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:162:36: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:173:11: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_generate_presplit_data.py:177:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:220:11: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_generate_presplit_data.py:225:33: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:227:31: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_generate_presplit_data.py:229:33: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_library_prediction.py:47:41: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_library_prediction.py:290:5: F401 'numpy as np' imported but unused
+./tests/tfscreen/simulate/test_library_prediction.py:439:16: E272 multiple spaces before keyword
./tests/tfscreen/simulate/test_run_simulation.py:3:1: F401 'os' imported but unused
./tests/tfscreen/simulate/test_run_simulation.py:5:1: F401 'unittest.mock.MagicMock' imported but unused
./tests/tfscreen/simulate/test_run_simulation.py:7:1: F401 'tfscreen.util' imported but unused
./tests/tfscreen/simulate/test_run_simulation.py:9:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_run_simulation.py:13:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:26:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:29:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:33:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:40:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:46:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:50:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:64:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:69:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:71:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:74:128: E501 line too long (128 > 127 characters)
-./tests/tfscreen/simulate/test_run_simulation.py:76:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:78:128: E501 line too long (148 > 127 characters)
-./tests/tfscreen/simulate/test_run_simulation.py:80:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:89:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:92:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:94:47: W291 trailing whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:97:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_run_simulation.py:99:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:103:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:106:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_run_simulation.py:109:5: F841 local variable 'results' is assigned to but never used
+./tests/tfscreen/simulate/test_run_simulation.py:27:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:30:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:34:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:41:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:47:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:51:1: W293 blank line contains whitespace
+./tests/tfscreen/simulate/test_run_simulation.py:55:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:65:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:70:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_run_simulation.py:82:128: E501 line too long (150 > 127 characters)
+./tests/tfscreen/simulate/test_run_simulation.py:101:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_run_simulation.py:110:1: W293 blank line contains whitespace
+./tests/tfscreen/simulate/test_run_simulation.py:113:1: W293 blank line contains whitespace
+./tests/tfscreen/simulate/test_run_simulation.py:116:5: F841 local variable 'results' is assigned to but never used
+./tests/tfscreen/simulate/test_run_simulation.py:117:1: W293 blank line contains whitespace
+./tests/tfscreen/simulate/test_sample_theta.py:260:21: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_sample_theta.py:277:27: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_sample_theta.py:282:21: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_sample_theta.py:297:21: E221 multiple spaces before operator
+./tests/tfscreen/simulate/test_sample_theta.py:338:21: E221 multiple spaces before operator
./tests/tfscreen/simulate/test_selection_experiment.py:4:1: F401 'numpy.ma' imported but unused
./tests/tfscreen/simulate/test_selection_experiment.py:6:1: F401 'scipy.stats.gmean' imported but unused
./tests/tfscreen/simulate/test_selection_experiment.py:9:1: F401 'unittest.mock.MagicMock' imported but unused
-./tests/tfscreen/simulate/test_selection_experiment.py:33:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_selection_experiment.py:38:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_selection_experiment.py:34:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_selection_experiment.py:39:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_selection_experiment.py:68:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_selection_experiment.py:77:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/simulate/test_selection_experiment.py:85:1: W293 blank line contains whitespace
@@ -5312,1116 +5446,1140 @@
./tests/tfscreen/simulate/test_selection_experiment.py:1265:1: W293 blank line contains whitespace
./tests/tfscreen/simulate/test_selection_experiment.py:1269:1: W293 blank line contains whitespace
./tests/tfscreen/simulate/test_selection_experiment.py:1272:1: W293 blank line contains whitespace
-./tests/tfscreen/simulate/test_selection_experiment.py:1274:73: W292 no newline at end of file
-./tests/tfscreen/simulate/test_thermo_to_growth.py:27:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_thermo_to_growth.py:31:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_thermo_to_growth.py:40:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/simulate/test_thermo_to_growth.py:95:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/simulate/test_thermo_to_growth.py:97:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/simulate/test_thermo_to_growth.py:103:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/simulate/test_thermo_to_growth.py:105:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/simulate/test_thermo_to_growth.py:172:14: E131 continuation line unaligned for hanging indent
-./tests/tfscreen/simulate/test_thermo_to_growth.py:199:44: E127 continuation line over-indented for visual indent
-./tests/tfscreen/simulate/test_thermo_to_growth.py:206:35: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:23:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:29:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:36:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:42:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:49:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:57:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:64:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:66:45: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:67:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:72:19: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:73:24: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:76:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:80:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:83:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:86:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:90:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:94:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:97:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_fixed.py:113:42: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:3:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:23:17: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:27:18: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:28:18: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:31:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:40:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:44:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:47:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:50:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:54:29: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:55:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:69:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:76:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:82:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:89:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:97:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:103:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:106:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:111:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:118:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:120:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:121:55: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:126:19: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:127:24: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:130:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:134:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:137:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:142:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:147:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:153:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py:356:6: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/activity/test_hierarchical_mut.py:36:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:3:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:26:18: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:27:18: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:30:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:39:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:43:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:46:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:49:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:53:29: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:54:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:68:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:75:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:81:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:85:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:87:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:91:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:94:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:97:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:100:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:104:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:110:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:113:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:117:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:124:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:126:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:127:55: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:132:19: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:133:24: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:136:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:140:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:143:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:147:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:152:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:158:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:174:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py:176:42: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/activity/test_horseshoe_mut.py:36:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:23:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:29:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:36:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:42:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:49:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:57:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:64:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:66:44: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:67:49: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:72:19: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:73:24: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:76:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:80:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:83:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:86:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:90:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:94:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:97:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py:113:41: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:4:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:78:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:82:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:120:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:121:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:183:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:197:34: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:282:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:283:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:314:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:352:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:353:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:370:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:383:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:386:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:389:21: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:393:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:409:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:410:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:411:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:412:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:415:21: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py:426:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/growth/test_linear.py:7:1: F401 'tfscreen.tfmodel.generative.components.growth.linear.LinearParams' imported but unused
-./tests/tfscreen/tfmodel/components/growth/test_linear.py:55:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/growth/test_linear.py:132:28: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/growth/test_linear.py:134:28: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/growth/test_power.py:3:1: F401 'numpyro.handlers.substitute' imported but unused
-./tests/tfscreen/tfmodel/components/growth/test_power.py:6:1: F401 'tfscreen.tfmodel.generative.components.growth.power.PowerParams' imported but unused
-./tests/tfscreen/tfmodel/components/growth/test_power.py:23:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/growth/test_saturation.py:6:1: F401 'tfscreen.tfmodel.generative.components.growth.saturation.SaturationParams' imported but unused
-./tests/tfscreen/tfmodel/components/growth/test_saturation.py:23:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/growth_noise/test_normal_kt.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/growth_noise/test_normal_kt.py:3:1: F401 'jax.numpy as jnp' imported but unused
-./tests/tfscreen/tfmodel/components/growth_noise/test_normal_kt.py:4:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/growth_noise/test_normal_kt.py:5:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/growth_noise/test_zero.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:8:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:12:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:19:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:23:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:28:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:30:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:35:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:38:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:42:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:47:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:49:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_instant.py:51:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/growth_transition/test_memory.py:99:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/growth_transition/test_memory.py:100:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/growth_transition/test_memory.py:101:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/growth_transition/test_memory.py:102:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/growth_transition/test_memory.py:103:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:5:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:6:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:38:35: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:74:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:76:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:77:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:80:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:107:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:109:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:110:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:113:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:139:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:141:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:142:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:145:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:181:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:183:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:184:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:187:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:221:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:223:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:224:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:227:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:232:77: E202 whitespace before ']'
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:455:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:471:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:508:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:509:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:555:12: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:556:12: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:574:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:576:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:577:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:618:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:621:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:669:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:671:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:672:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:675:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:715:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:718:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:723:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:748:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:750:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:752:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:760:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:773:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:775:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:777:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:786:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:798:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:799:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:802:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:816:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:817:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:835:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:837:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:838:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:899:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:901:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:902:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:958:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:962:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:994:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:1015:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:1136:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py:1254:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:5:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:6:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:115:77: E202 whitespace before ']'
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:144:32: E131 continuation line unaligned for hanging indent
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:146:32: E131 continuation line unaligned for hanging indent
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:149:38: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:201:12: F541 f-string is missing placeholders
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:208:12: F541 f-string is missing placeholders
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:254:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:256:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:258:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:267:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:269:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:288:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:290:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:293:7: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:294:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:306:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:310:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:311:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:312:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:323:7: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:324:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:338:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:343:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:363:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:370:7: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:385:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:393:7: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:404:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:412:7: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:425:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:426:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:442:7: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:462:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:474:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:476:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:488:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:490:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:502:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:514:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:516:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:526:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:528:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:538:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:556:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:582:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:616:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:634:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py:653:9: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:20:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:24:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:30:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:35:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:41:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:45:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:47:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:51:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:54:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:60:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:66:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:69:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:77:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:85:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:88:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:100:19: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:101:25: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:104:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:109:30: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:110:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:114:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:118:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:121:65: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:122:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:125:64: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:126:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:131:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:137:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:158:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:161:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:165:58: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:166:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:170:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:175:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:183:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:192:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:197:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:201:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:206:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:211:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:215:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_beta.py:217:58: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/noise/test_logit_normal.py:4:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:3:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:5:1: F401 'collections.namedtuple' imported but unused
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:19:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:22:18: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:26:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:32:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:39:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:44:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:51:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:54:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:56:42: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:57:51: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:61:37: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:63:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:66:19: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:67:28: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:70:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:75:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:82:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:85:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:87:35: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:88:44: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/noise/test_zero.py:93:46: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/sample_offset/test_normal.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/sample_offset/test_normal.py:5:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/sample_offset/test_normal.py:6:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/sample_offset/test_zero.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/sample_offset/test_zero.py:2:1: F401 'jax.numpy as jnp' imported but unused
-./tests/tfscreen/tfmodel/components/test_pinning.py:4:1: F401 'numpyro as pyro' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test__simple.py:6:1: F401 'tfscreen.tfmodel.generative.components.theta._simple._THETA_EPS' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:5:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:6:1: F401 'numpyro.handlers.trace' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:10:1: F401 'tfscreen.tfmodel.generative.components.theta.categorical_geno.ThetaParam' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:38:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:46:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:51:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:55:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:69:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:77:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:86:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:91:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:96:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:102:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:107:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:119:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:125:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:130:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:146:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:153:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:158:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:167:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:172:21: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:174:34: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:175:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:176:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:180:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:186:37: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:188:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:190:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:193:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:200:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:220:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py:379:44: W292 no newline at end of file
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:4:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:5:1: F401 'numpyro.handlers.trace' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:36:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:171:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:174:36: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:175:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:200:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:243:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:247:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:248:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:273:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:274:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:280:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:280:64: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:299:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:315:20: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:324:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:326:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:350:38: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:352:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:353:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:472:51: E222 multiple spaces after operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_geno.py:487:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:9:1: F401 'tfscreen.tfmodel.generative.components.theta.hill_mut._population_moments' imported but unused
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:49:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:55:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:62:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:119:15: F841 local variable 'G' is assigned to but never used
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:130:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:134:15: F841 local variable 'G' is assigned to but never used
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:134:18: F841 local variable 'P' is assigned to but never used
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:240:9: F841 local variable 'M' is assigned to but never used
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:248:62: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_hill_mut.py:447:56: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:25:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:30:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:155:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:156:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:157:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:158:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:159:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:160:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:161:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:162:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:164:36: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:166:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:167:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:204:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:205:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py:301:19: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:41:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:58:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:195:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:248:6: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:280:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:315:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:332:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:333:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:346:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:347:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:424:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:425:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:426:45: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:427:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:428:34: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:429:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:430:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:443:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:485:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:596:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:615:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:615:28: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:615:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:616:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:617:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:617:28: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:617:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:618:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:618:28: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:618:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:685:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:693:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:43:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:97:45: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:138:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:195:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:197:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:223:38: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:245:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:266:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:277:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:308:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:327:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:391:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:473:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:478:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:509:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:528:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:528:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:528:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:529:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:530:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:530:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:530:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:531:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:531:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:531:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:588:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:591:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:595:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:599:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:602:15: E222 multiple spaces after operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:754:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:755:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:756:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:757:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:759:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:42:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:113:45: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:165:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:249:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:251:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:299:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:305:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:306:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:338:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:349:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:380:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:418:26: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:426:25: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:457:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:463:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:464:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:543:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:545:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:571:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:590:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:590:28: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:590:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:591:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:592:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:592:28: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:592:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:593:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:593:28: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:593:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:647:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:5:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:21:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:32:14: E222 multiple spaces after operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:258:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:26:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:35:6: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:217:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:217:33: E201 whitespace after '['
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:219:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:49:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:66:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:203:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:275:6: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:276:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:277:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:278:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:292:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:299:45: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:306:46: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:321:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:360:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:361:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:378:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:379:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:394:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:395:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:421:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:422:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:509:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:520:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:521:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:523:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:524:34: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:525:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:526:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:540:39: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:541:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:542:45: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:543:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:544:34: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:545:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:546:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:559:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:606:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:607:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:617:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:709:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:745:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:764:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:764:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:764:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:765:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:766:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:766:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:766:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:767:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:767:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:767:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:795:24: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:803:24: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:812:26: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:820:26: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:830:28: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:844:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:852:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:53:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:107:45: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:147:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:204:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:206:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:234:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:239:38: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:272:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:285:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:294:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:305:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:351:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:380:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:399:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:433:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:438:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:531:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:536:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:567:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:586:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:586:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:586:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:587:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:588:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:588:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:588:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:589:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:589:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:589:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:616:24: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:648:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:651:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:655:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:658:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:659:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:662:15: E222 multiple spaces after operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:667:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:694:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:695:29: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:696:29: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:697:29: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:698:29: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:699:29: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:820:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:49:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:121:45: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:163:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:230:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:232:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:263:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:287:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:288:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:298:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:299:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:320:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:329:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:340:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:386:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:427:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:430:26: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:447:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:453:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:454:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:544:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:546:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:572:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:591:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:591:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:591:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:592:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:593:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:593:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:593:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:594:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:594:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:594:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:621:24: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:646:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:28:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:39:6: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:43:8: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:310:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:310:33: E201 whitespace after '['
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:312:32: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:465:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:486:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:514:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:516:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:47:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:48:25: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:61:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:68:1: E305 expected 2 blank lines after class or function definition, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:68:6: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:282:38: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:283:43: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:284:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:343:38: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:344:43: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:345:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:346:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:347:34: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:348:46: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:349:46: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:361:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:362:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:363:42: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:364:42: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:365:42: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:366:41: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:376:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:47:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:98:45: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:123:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:134:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:265:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:274:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:292:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:308:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:383:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:448:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:467:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:467:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:467:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:468:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:469:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:469:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:469:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:470:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:470:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:470:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:504:26: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:9:1: F401 'functools.partial' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:40:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:211:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:212:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:213:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:214:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:220:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:221:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:237:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:270:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:271:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:272:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:273:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:278:27: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:279:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:280:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:281:32: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:282:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:283:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:284:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:318:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:325:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:326:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:343:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:381:38: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:383:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:491:26: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:503:28: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:513:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:5:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:21:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:33:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:51:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:52:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:53:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:54:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:58:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:71:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:72:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:73:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:74:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:78:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:125:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:244:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:66:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:73:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:82:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:89:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:98:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:101:19: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:259:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:263:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:47:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:48:25: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:61:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:68:1: E305 expected 2 blank lines after class or function definition, found 1
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:68:6: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:313:38: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:314:43: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:315:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:326:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:327:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:328:43: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:348:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:349:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:350:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:351:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:353:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:354:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:400:38: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:401:43: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:402:44: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:403:35: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:404:34: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:405:46: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:406:46: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:417:37: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:418:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:419:43: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:431:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:432:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:433:42: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:434:42: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:435:42: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:436:41: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:438:41: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:448:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:450:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:506:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:549:46: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:550:45: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:551:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:554:48: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:555:47: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:556:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:558:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:564:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:47:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:100:45: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:215:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:220:38: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:229:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:253:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:265:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:266:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:275:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:286:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:310:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:311:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:323:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:340:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:365:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:432:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:433:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:461:37: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:494:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:495:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:496:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:499:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:530:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:549:17: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:549:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:549:40: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:550:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:551:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:551:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:551:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:552:19: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:552:29: E701 multiple statements on one line (colon)
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:552:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:9:1: F401 'functools.partial' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:39:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:208:41: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:218:42: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:229:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:230:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:231:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:232:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:233:34: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:244:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:245:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:261:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:272:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:324:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:326:33: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:343:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:350:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:351:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:400:28: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:448:38: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:449:40: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:450:38: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:451:40: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:452:39: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:460:37: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:488:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:489:32: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:510:30: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:582:26: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:594:28: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:604:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:67:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:75:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:84:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:91:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:100:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:103:19: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:109:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:165:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:171:12: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:216:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:310:12: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:319:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:323:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:62:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:63:22: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:150:19: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:168:9: F841 local variable 'sdata' is assigned to but never used
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:177:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:178:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_features.py:179:33: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:6:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:7:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:30:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:37:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:38:31: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:39:30: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:43:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:49:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:61:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:71:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:76:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:82:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py:89:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:39:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:41:21: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:130:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:147:21: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:166:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:167:11: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:172:31: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:229:49: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:236:49: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:241:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:244:55: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:272:13: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_io.py:377:49: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:3:1: F401 'jax' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:7:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:9:1: F401 'tfscreen.tfmodel.generative.components.theta.thermo.nn._DEFAULT_HIDDEN_SIZE' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:24:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:26:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:39:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:41:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:54:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:72:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:74:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:77:9: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:95:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py:97:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:6:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:7:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:29:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:36:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:43:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:51:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:75:21: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:84:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py:91:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/components/theta_rescale/test_logit.py:2:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:3:1: F401 'numpy as np' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:5:1: F401 'numpyro.distributions as dist' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:6:1: F401 'unittest.mock.patch' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:16:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:22:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:24:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:29:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:33:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:37:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:39:32: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:40:36: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:41:31: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:53:26: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:54:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:56:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:58:24: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:59:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:63:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:65:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:68:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:71:26: W291 trailing whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:73:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:75:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:79:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:82:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:84:35: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:85:38: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:86:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:88:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:93:55: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:103:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:104:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:109:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:126:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:131:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:135:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:141:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:148:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:161:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:166:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:169:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test__congression.py:185:1: W391 blank line at end of file
-./tests/tfscreen/tfmodel/components/transformation/test_empirical.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test_logit_norm.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test_logit_norm.py:3:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:1:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:4:1: F401 'collections.namedtuple' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:7:1: F401 'tfscreen.tfmodel.data_class.GrowthData' imported but unused
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:9:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:15:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:21:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:26:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:31:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:36:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:39:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/components/transformation/test_single.py:43:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:3:1: F401 'numpyro' imported but unused
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:25:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:32:22: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:33:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:36:51: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:37:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:38:26: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:52:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:59:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:62:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:69:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:73:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:75:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:79:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:83:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:87:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:92:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:95:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:97:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:105:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:108:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:115:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_binding.py:117:41: W292 no newline at end of file
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:22:20: W291 trailing whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:28:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:32:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:42:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:45:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:47:76: W291 trailing whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:49:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:52:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:55:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:58:31: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:58:33: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:58:35: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:58:37: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:58:39: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:58:41: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:59:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:75:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:99:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:107:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:112:56: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:113:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:118:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:125:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:127:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:133:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:137:24: W291 trailing whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:138:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:146:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:151:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:154:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:156:35: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:156:37: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:156:39: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:156:41: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:156:43: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:156:45: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:162:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:168:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:171:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:173:23: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:173:25: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:173:27: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:173:29: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:173:31: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:173:33: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:174:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:176:23: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:176:25: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:176:27: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:176:29: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:176:31: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:176:33: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:178:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:184:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:192:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:197:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/observe/test_observe_growth.py:259:57: W292 no newline at end of file
+./tests/tfscreen/simulate/test_selection_experiment.py:1313:27: W292 no newline at end of file
+./tests/tfscreen/simulate/test_simulate_cli.py:3:1: F401 'os' imported but unused
+./tests/tfscreen/simulate/test_simulate_cli.py:247:5: F401 'tfscreen.simulate.scripts.simulate_cli as cli_mod' imported but unused
+./tests/tfscreen/simulate/test_thermo_to_growth.py:28:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_thermo_to_growth.py:32:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_thermo_to_growth.py:41:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/simulate/test_thermo_to_growth.py:96:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:98:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:104:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:106:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:195:14: E131 continuation line unaligned for hanging indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:222:44: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:229:35: E127 continuation line over-indented for visual indent
+./tests/tfscreen/simulate/test_thermo_to_growth.py:1437:1: E402 module level import not at top of file
+./tests/tfscreen/simulate/test_thermo_to_growth.py:1525:9: F841 local variable 'G_out' is assigned to but never used
+./tests/tfscreen/tfmodel/analysis/test_error_calibration.py:313:9: F401 'matplotlib.axes.Axes' imported but unused
+./tests/tfscreen/tfmodel/analysis/test_error_calibration.py:355:9: F401 'matplotlib.axes.Axes' imported but unused
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:23:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:29:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:36:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:42:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:49:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:54:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:57:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:64:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:66:45: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:67:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:72:19: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:73:24: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:76:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:80:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:83:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:86:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:90:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:94:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:97:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py:113:42: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:3:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:23:17: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:27:18: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:28:18: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:31:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:40:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:44:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:47:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:50:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:54:29: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:55:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:69:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:76:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:82:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:89:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:97:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:103:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:106:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:111:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:118:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:120:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:121:55: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:126:19: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:127:24: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:130:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:134:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:137:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:142:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:147:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:153:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py:356:6: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_mut.py:36:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:3:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:26:18: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:27:18: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:30:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:39:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:43:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:46:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:49:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:53:29: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:54:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:68:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:75:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:81:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:85:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:87:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:91:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:94:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:97:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:100:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:104:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:110:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:113:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:117:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:124:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:126:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:127:55: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:132:19: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:133:24: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:136:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:140:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:143:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:147:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:152:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:158:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:174:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py:176:42: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_mut.py:36:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:23:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:29:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:36:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:42:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:49:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:54:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:57:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:64:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:66:44: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:67:49: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:72:19: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:73:24: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:76:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:80:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:83:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:86:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:90:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:94:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:97:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py:113:41: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:4:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:78:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:82:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:120:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:121:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:183:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:197:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:282:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:283:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:314:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:352:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:353:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:370:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:383:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:386:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:389:21: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:393:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:409:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:410:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:411:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:412:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:415:21: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py:426:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/growth/test_linear.py:7:1: F401 'tfscreen.tfmodel.generative.components.growth.linear.LinearParams' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth/test_linear.py:56:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/growth/test_linear.py:138:28: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/growth/test_linear.py:140:28: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/growth/test_power.py:3:1: F401 'numpyro.handlers.substitute' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth/test_power.py:6:1: F401 'tfscreen.tfmodel.generative.components.growth.power.PowerParams' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth/test_power.py:23:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/growth/test_saturation.py:6:1: F401 'tfscreen.tfmodel.generative.components.growth.saturation.SaturationParams' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth/test_saturation.py:23:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/growth_noise/test_normal_kt.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth_noise/test_normal_kt.py:3:1: F401 'jax.numpy as jnp' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth_noise/test_normal_kt.py:4:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth_noise/test_normal_kt.py:5:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth_noise/test_zero.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:8:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:12:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:19:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:23:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:28:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:30:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:35:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:38:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:42:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:47:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:49:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py:51:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py:99:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py:100:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py:101:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py:102:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py:103:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:5:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:6:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:38:35: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:74:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:76:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:77:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:80:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:107:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:109:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:110:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:113:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:139:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:141:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:142:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:145:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:181:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:183:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:184:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:187:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:221:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:223:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:224:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:227:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:232:77: E202 whitespace before ']'
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:455:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:471:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:508:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:509:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:555:12: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:556:12: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:574:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:576:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:577:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:618:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:621:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:669:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:671:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:672:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:675:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:715:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:718:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:723:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:748:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:750:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:752:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:760:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:773:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:775:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:777:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:786:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:798:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:799:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:802:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:816:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:817:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:835:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:837:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:838:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:899:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:901:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:902:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:958:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:962:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:994:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:1015:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:1136:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py:1254:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:5:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:6:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:115:77: E202 whitespace before ']'
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:144:32: E131 continuation line unaligned for hanging indent
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:146:32: E131 continuation line unaligned for hanging indent
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:149:38: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:201:12: F541 f-string is missing placeholders
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:208:12: F541 f-string is missing placeholders
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:254:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:256:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:258:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:267:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:269:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:288:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:290:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:293:7: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:294:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:306:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:310:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:311:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:312:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:323:7: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:324:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:338:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:343:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:363:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:370:7: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:385:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:393:7: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:404:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:412:7: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:425:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:426:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:442:7: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:462:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:474:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:476:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:488:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:490:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:502:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:514:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:516:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:526:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:528:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:538:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:556:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:582:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:616:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:634:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py:653:9: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:20:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:24:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:30:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:35:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:41:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:45:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:47:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:51:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:54:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:60:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:66:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:69:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:77:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:85:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:88:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:100:19: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:101:25: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:104:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:109:30: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:110:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:114:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:118:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:121:65: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:122:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:125:64: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:126:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:131:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:137:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:158:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:161:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:165:58: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:166:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:170:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:175:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:183:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:192:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:197:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:201:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:206:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:211:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:215:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_beta.py:217:58: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/noise/test_logit_normal.py:4:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:3:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:5:1: F401 'collections.namedtuple' imported but unused
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:19:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:22:18: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:26:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:32:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:39:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:44:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:51:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:54:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:56:42: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:57:51: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:61:37: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:63:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:66:19: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:67:28: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:70:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:75:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:82:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:85:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:87:35: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:88:44: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/noise/test_zero.py:93:46: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/sample_offset/test_normal.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/sample_offset/test_normal.py:5:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/sample_offset/test_normal.py:6:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/sample_offset/test_zero.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/sample_offset/test_zero.py:2:1: F401 'jax.numpy as jnp' imported but unused
+./tests/tfscreen/tfmodel/generative/components/test_pinning.py:4:1: F401 'numpyro as pyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test__simple.py:6:1: F401 'tfscreen.tfmodel.generative.components.theta._simple._THETA_EPS' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test__simple.py:447:34: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:5:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:6:1: F401 'numpyro.handlers.trace' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:10:1: F401 'tfscreen.tfmodel.generative.components.theta.categorical_geno.ThetaParam' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:38:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:46:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:51:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:55:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:69:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:77:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:86:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:91:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:96:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:102:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:107:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:119:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:125:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:130:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:146:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:153:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:158:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:167:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:172:21: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:174:34: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:175:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:176:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:180:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:186:37: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:188:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:190:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:193:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:200:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:220:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py:379:44: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:4:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:5:1: F401 'numpyro.handlers.trace' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:41:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:176:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:179:36: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:180:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:205:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:248:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:252:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:253:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:278:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:279:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:285:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:285:64: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:304:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:320:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:329:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:331:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:355:38: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:357:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:358:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:477:51: E222 multiple spaces after operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:492:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:537:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:540:36: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py:541:26: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:12:1: F401 'tfscreen.tfmodel.generative.components.theta.hill_mut._population_moments' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:55:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:61:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:68:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:125:15: F841 local variable 'G' is assigned to but never used
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:136:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:140:15: F841 local variable 'G' is assigned to but never used
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:140:18: F841 local variable 'P' is assigned to but never used
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:246:9: F841 local variable 'M' is assigned to but never used
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:254:62: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:453:56: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:480:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:482:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:492:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:494:26: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:509:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:512:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:513:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:528:21: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:530:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:530:52: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:719:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:722:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:730:36: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py:730:70: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:25:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:30:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:155:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:156:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:157:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:158:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:159:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:160:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:161:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:162:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:164:36: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:166:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:167:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:204:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:205:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py:301:19: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:41:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:58:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:195:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:248:6: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:280:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:315:27: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:332:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:333:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:346:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:347:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:424:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:425:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:426:45: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:427:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:428:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:429:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:430:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:443:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:485:27: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:596:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:615:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:615:28: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:615:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:616:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:617:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:617:28: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:617:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:618:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:618:28: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:618:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:685:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py:693:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:43:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:97:45: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:138:23: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:195:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:197:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:223:38: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:245:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:266:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:277:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:308:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:327:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:391:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:473:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:478:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:509:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:528:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:528:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:528:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:529:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:530:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:530:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:530:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:531:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:531:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:531:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:588:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:591:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:595:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:599:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:602:15: E222 multiple spaces after operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:754:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:755:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:756:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:757:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py:759:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:42:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:113:45: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:165:23: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:249:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:251:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:299:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:305:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:306:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:338:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:349:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:380:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:418:26: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:426:25: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:457:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:463:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:464:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:543:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:545:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:571:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:590:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:590:28: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:590:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:591:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:592:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:592:28: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:592:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:593:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:593:28: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:593:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py:647:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:5:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:21:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:32:14: E222 multiple spaces after operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py:258:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:26:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:35:6: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:217:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:217:33: E201 whitespace after '['
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py:219:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:49:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:66:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:203:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:275:6: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:276:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:277:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:278:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:292:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:299:45: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:306:46: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:321:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:360:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:361:27: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:378:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:379:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:394:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:395:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:421:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:422:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:509:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:520:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:521:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:523:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:524:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:525:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:526:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:540:39: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:541:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:542:45: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:543:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:544:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:545:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:546:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:559:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:606:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:607:27: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:617:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:709:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:745:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:764:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:764:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:764:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:765:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:766:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:766:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:766:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:767:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:767:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:767:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:795:24: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:803:24: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:812:26: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:820:26: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:830:28: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:844:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py:852:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:53:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:107:45: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:147:23: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:204:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:206:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:234:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:239:38: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:272:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:285:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:294:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:305:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:351:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:380:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:399:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:433:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:438:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:531:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:536:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:567:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:586:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:586:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:586:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:587:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:588:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:588:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:588:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:589:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:589:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:589:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:616:24: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:648:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:651:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:655:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:658:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:659:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:662:15: E222 multiple spaces after operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:667:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:694:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:695:29: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:696:29: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:697:29: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:698:29: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:699:29: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py:820:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:49:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:121:45: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:163:23: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:230:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:232:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:263:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:287:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:288:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:298:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:299:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:320:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:329:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:340:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:386:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:427:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:430:26: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:447:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:453:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:454:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:544:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:546:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:572:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:591:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:591:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:591:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:592:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:593:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:593:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:593:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:594:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:594:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:594:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:621:24: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py:646:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:28:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:39:6: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:43:8: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:310:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:310:33: E201 whitespace after '['
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:312:32: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:465:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:486:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:514:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py:516:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:47:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:48:25: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:54:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:61:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:68:1: E305 expected 2 blank lines after class or function definition, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:68:6: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:282:38: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:283:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:284:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:343:38: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:344:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:345:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:346:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:347:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:348:46: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:349:46: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:361:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:362:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:363:42: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:364:42: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:365:42: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:366:41: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py:376:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:47:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:98:45: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:123:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:134:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:265:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:274:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:292:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:308:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:383:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:448:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:467:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:467:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:467:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:468:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:469:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:469:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:469:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:470:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:470:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:470:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py:504:26: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:9:1: F401 'functools.partial' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:40:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:211:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:212:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:213:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:214:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:220:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:221:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:237:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:270:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:271:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:272:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:273:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:278:27: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:279:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:280:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:281:32: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:282:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:283:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:284:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:318:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:325:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:326:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:343:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:381:38: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:383:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:491:26: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:503:28: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py:513:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:5:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:21:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:33:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:51:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:52:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:53:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:54:27: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:58:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:71:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:72:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:73:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:74:27: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:78:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:125:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py:244:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:66:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:73:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:82:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:89:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:98:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:101:19: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:259:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py:263:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:47:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:48:25: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:54:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:61:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:68:1: E305 expected 2 blank lines after class or function definition, found 1
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:68:6: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:313:38: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:314:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:315:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:326:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:327:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:328:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:348:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:349:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:350:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:351:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:353:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:354:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:400:38: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:401:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:402:44: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:403:35: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:404:34: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:405:46: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:406:46: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:417:37: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:418:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:419:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:431:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:432:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:433:42: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:434:42: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:435:42: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:436:41: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:438:41: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:448:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:450:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:506:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:549:46: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:550:45: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:551:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:554:48: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:555:47: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:556:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:558:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py:564:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:47:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:100:45: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:215:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:220:38: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:229:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:253:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:265:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:266:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:275:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:286:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:310:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:311:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:323:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:340:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:365:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:432:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:433:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:461:37: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:494:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:495:23: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:496:23: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:499:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:530:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:549:17: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:549:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:549:40: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:550:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:551:18: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:551:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:551:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:552:19: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:552:29: E701 multiple statements on one line (colon)
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py:552:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:9:1: F401 'functools.partial' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:39:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:208:41: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:218:42: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:229:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:230:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:231:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:232:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:233:34: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:244:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:245:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:261:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:272:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:324:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:326:33: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:343:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:350:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:351:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:400:28: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:448:38: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:449:40: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:450:38: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:451:40: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:452:39: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:460:37: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:488:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:489:32: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:510:30: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:582:26: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:594:28: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py:604:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:67:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:75:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:84:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:91:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:100:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:103:19: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:109:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:165:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:171:12: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:216:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:310:12: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:319:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py:323:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:62:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:63:22: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:150:19: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:168:9: F841 local variable 'sdata' is assigned to but never used
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:177:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:178:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py:179:33: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:6:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:7:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:30:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:37:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:38:31: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:39:30: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:43:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:49:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:61:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:71:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:76:15: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:82:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py:89:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:39:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:41:21: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:130:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:147:21: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:166:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:167:11: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:172:31: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:229:49: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:236:49: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:241:10: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:244:55: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:272:13: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py:377:49: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:3:1: F401 'jax' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:7:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:9:1: F401 'tfscreen.tfmodel.generative.components.theta.thermo.nn._DEFAULT_HIDDEN_SIZE' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:24:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:26:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:39:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:41:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:54:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:72:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:74:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:77:9: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:95:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py:97:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:6:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:7:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:29:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:36:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:43:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:51:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:75:21: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:84:16: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py:91:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/generative/components/theta_rescale/test_logit.py:2:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:3:1: F401 'numpy as np' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:5:1: F401 'numpyro.distributions as dist' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:6:1: F401 'unittest.mock.patch' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:16:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:22:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:24:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:29:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:33:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:37:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:39:32: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:40:36: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:41:47: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:53:26: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:54:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:56:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:58:24: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:59:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:63:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:65:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:68:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:71:26: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:73:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:75:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:79:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:82:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:84:35: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:85:38: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:86:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:88:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:93:55: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:103:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:104:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:109:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:126:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:131:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:135:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:141:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:148:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:161:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:166:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:169:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py:185:1: W391 blank line at end of file
+./tests/tfscreen/tfmodel/generative/components/transformation/test_empirical.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test_logit_norm.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test_logit_norm.py:3:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:1:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:4:1: F401 'collections.namedtuple' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:7:1: F401 'tfscreen.tfmodel.data_class.GrowthData' imported but unused
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:9:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:15:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:21:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:26:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:31:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:36:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:39:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/components/transformation/test_single.py:43:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:3:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:25:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:32:22: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:33:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:36:51: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:37:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:38:26: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:52:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:59:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:62:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:69:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:73:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:75:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:79:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:83:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:87:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:92:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:95:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:97:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:105:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:108:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:115:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py:117:41: W292 no newline at end of file
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:22:20: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:28:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:32:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:42:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:45:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:47:76: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:49:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:52:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:55:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:58:31: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:58:33: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:58:35: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:58:37: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:58:39: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:58:41: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:59:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:75:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:99:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:107:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:112:56: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:113:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:118:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:125:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:127:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:133:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:137:24: W291 trailing whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:138:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:146:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:151:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:154:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:156:35: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:156:37: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:156:39: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:156:41: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:156:43: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:156:45: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:162:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:168:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:171:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:173:23: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:173:25: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:173:27: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:173:29: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:173:31: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:173:33: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:174:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:176:23: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:176:25: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:176:27: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:176:29: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:176:31: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:176:33: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:178:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:184:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:192:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:197:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py:259:57: W292 no newline at end of file
./tests/tfscreen/tfmodel/scripts/test_diagnose_nan_cli.py:77:64: E203 whitespace before ','
./tests/tfscreen/tfmodel/scripts/test_extract_params_cli.py:14:1: F401 'numpyro.infer.svi.SVIState' imported but unused
./tests/tfscreen/tfmodel/scripts/test_fit_model_cli.py:5:1: F401 'dill' imported but unused
@@ -6433,48 +6591,52 @@
./tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py:4:1: F401 'os' imported but unused
./tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py:7:1: F401 'numpy as np' imported but unused
./tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py:8:1: F401 'unittest.mock.call' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_prior_predictive_cli.py:3:1: F401 'pytest' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:11:1: F401 'dataclasses' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:13:1: F401 'shutil' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:14:1: F401 'sys' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:26:1: F401 'tfscreen.tfmodel.scripts.run_prefit_calibration_cli._identify_field_mapping' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:295:1: E303 too many blank lines (4)
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:299:1: E302 expected 2 blank lines, found 4
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1003:67: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1039:14: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1058:14: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1073:14: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1093:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1099:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1100:15: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1103:17: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1104:13: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1105:13: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1106:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1110:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1114:14: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1118:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1119:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1120:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1121:18: E272 multiple spaces before keyword
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1148:24: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1150:10: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1167:23: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1168:27: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1175:16: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1176:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1640:12: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1666:63: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1705:38: E702 multiple statements on one line (semicolon)
-./tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py:1724:20: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:11:1: F401 'dataclasses' imported but unused
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:13:1: F401 'shutil' imported but unused
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:14:1: F401 'sys' imported but unused
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:21:1: F401 'jax' imported but unused
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:26:1: F401 'tfscreen.tfmodel.scripts.prefit_calibration_cli._identify_field_mapping' imported but unused
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:733:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:734:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:735:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:736:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:737:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:738:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:739:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:740:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:741:29: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1117:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1123:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1129:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1136:47: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1144:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1150:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1158:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1164:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1170:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1177:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1184:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1198:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1212:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1224:128: E501 line too long (136 > 127 characters)
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1309:14: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1328:14: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py:1343:14: E127 continuation line over-indented for visual indent
./tests/tfscreen/tfmodel/scripts/test_sample_posterior_cli.py:5:1: F401 'unittest.mock.call' imported but unused
-./tests/tfscreen/tfmodel/scripts/test_subset_genotypes.py:84:31: E741 ambiguous variable name 'l'
-./tests/tfscreen/tfmodel/scripts/test_subset_genotypes.py:598:28: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:15:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/scripts/test_sample_prior_cli.py:3:1: F401 'pytest' imported but unused
+./tests/tfscreen/tfmodel/scripts/test_subset_genotypes_cli.py:84:31: E741 ambiguous variable name 'l'
+./tests/tfscreen/tfmodel/scripts/test_subset_genotypes_cli.py:598:28: E127 continuation line over-indented for visual indent
./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:16:1: E402 module level import not at top of file
./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:17:1: E402 module level import not at top of file
./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:18:1: E402 module level import not at top of file
-./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:20:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:19:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:21:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:31:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:969:18: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:972:17: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:1237:1: E402 module level import not at top of file
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:1419:14: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py:1751:21: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/scripts/test_summarize_grid_cli.py:79:9: F841 local variable 'data' is assigned to but never used
./tests/tfscreen/tfmodel/scripts/test_summarize_grid_cli.py:183:91: E712 comparison to True should be 'if cond is True:' or 'if cond:'
./tests/tfscreen/tfmodel/test_batch.py:3:1: F401 'numpy as np' imported but unused
@@ -6557,83 +6719,86 @@
./tests/tfscreen/tfmodel/test_checkpoint_protection.py:42:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_checkpoint_protection.py:62:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_checkpoint_protection.py:78:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:20:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:24:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:27:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:32:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:42:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:52:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:55:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:58:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:63:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:69:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:78:39: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:79:40: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:80:42: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:82:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:84:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:88:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:91:45: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:92:48: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:93:44: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:94:41: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:96:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:100:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:109:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:117:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:121:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:123:71: E203 whitespace before ','
-./tests/tfscreen/tfmodel/test_configure_and_run.py:124:1: E128 continuation line under-indented for visual indent
-./tests/tfscreen/tfmodel/test_configure_and_run.py:125:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:127:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:132:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:134:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:137:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:142:30: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:142:40: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:142:55: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:142:74: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:145:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:148:30: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:148:40: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:148:55: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:148:74: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:154:34: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:154:49: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:156:14: E111 indentation is not a multiple of 4
-./tests/tfscreen/tfmodel/test_configure_and_run.py:156:14: E117 over-indented
-./tests/tfscreen/tfmodel/test_configure_and_run.py:156:31: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:156:41: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:156:56: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:156:75: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:162:34: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:162:54: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:14: E111 indentation is not a multiple of 4
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:14: E117 over-indented
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:31: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:41: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:56: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:75: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:164:128: E501 line too long (142 > 127 characters)
-./tests/tfscreen/tfmodel/test_configure_and_run.py:165:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_configure_and_run.py:168:66: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:14: E111 indentation is not a multiple of 4
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:14: E117 over-indented
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:31: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:41: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:56: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:75: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:174:128: E501 line too long (150 > 127 characters)
-./tests/tfscreen/tfmodel/test_configure_and_run.py:178:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:183:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:185:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/test_configure_and_run.py:188:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/test_configure_and_run.py:316:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:331:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_configure_and_run.py:339:114: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:342:9: F841 local variable 'mock_ri_class' is assigned to but never used
-./tests/tfscreen/tfmodel/test_configure_and_run.py:345:44: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_configure_and_run.py:361:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:21:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:25:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:28:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:33:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:43:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:53:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:56:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:59:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:64:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:70:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:79:39: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:80:40: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:81:42: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:83:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:85:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:89:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:92:45: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:93:48: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:94:44: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:95:41: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:97:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:101:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:110:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:118:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:122:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:124:71: E203 whitespace before ','
+./tests/tfscreen/tfmodel/test_configure_and_run.py:125:1: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/test_configure_and_run.py:126:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:133:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:135:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:138:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:143:30: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:143:40: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:143:55: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:143:74: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:146:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:149:30: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:149:40: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:149:55: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:149:74: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:155:34: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:155:49: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:157:14: E111 indentation is not a multiple of 4
+./tests/tfscreen/tfmodel/test_configure_and_run.py:157:14: E117 over-indented
+./tests/tfscreen/tfmodel/test_configure_and_run.py:157:31: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:157:41: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:157:56: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:157:75: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:163:34: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:163:54: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:14: E111 indentation is not a multiple of 4
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:14: E117 over-indented
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:31: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:41: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:56: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:75: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:165:128: E501 line too long (142 > 127 characters)
+./tests/tfscreen/tfmodel/test_configure_and_run.py:166:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_configure_and_run.py:169:76: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:14: E111 indentation is not a multiple of 4
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:14: E117 over-indented
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:31: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:41: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:56: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:75: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:175:128: E501 line too long (150 > 127 characters)
+./tests/tfscreen/tfmodel/test_configure_and_run.py:187:30: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:187:40: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:187:55: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:187:74: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:242:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:244:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/test_configure_and_run.py:247:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/test_configure_and_run.py:375:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:390:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_configure_and_run.py:398:124: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:398:128: E501 line too long (130 > 127 characters)
+./tests/tfscreen/tfmodel/test_configure_and_run.py:401:9: F841 local variable 'mock_ri_class' is assigned to but never used
+./tests/tfscreen/tfmodel/test_configure_and_run.py:404:44: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_configure_and_run.py:420:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_congression_mask.py:1:1: F401 'pytest' imported but unused
./tests/tfscreen/tfmodel/test_congression_mask.py:5:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_congression_mask.py:7:1: W293 blank line contains whitespace
@@ -6662,15 +6827,12 @@
./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:50:47: E261 at least two spaces before inline comment
./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:51:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:56:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:59:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:64:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:71:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:78:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:84:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:89:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:94:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:99:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:120:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:81:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:86:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:91:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:96:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extract_growth_predictions.py:117:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_extract_parameters_congression.py:8:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_extract_parameters_congression.py:30:128: E501 line too long (130 > 127 characters)
./tests/tfscreen/tfmodel/test_extract_parameters_congression.py:37:1: E302 expected 2 blank lines, found 1
@@ -6700,8 +6862,6 @@
./tests/tfscreen/tfmodel/test_extract_theta_curves.py:38:71: E261 at least two spaces before inline comment
./tests/tfscreen/tfmodel/test_extract_theta_curves.py:43:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_extract_theta_curves.py:46:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extract_theta_curves.py:53:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extract_theta_curves.py:57:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_extract_theta_curves.py:65:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_extract_theta_curves.py:71:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_extract_theta_curves.py:73:1: W293 blank line contains whitespace
@@ -6733,7 +6893,7 @@
./tests/tfscreen/tfmodel/test_extraction_new_models.py:209:17: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_extraction_new_models.py:211:13: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_extraction_new_models.py:399:18: E221 multiple spaces before operator
-./tests/tfscreen/tfmodel/test_extraction_new_models.py:407:43: E221 multiple spaces before operator
+./tests/tfscreen/tfmodel/test_extraction_new_models.py:407:41: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_extraction_new_models.py:479:12: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_extraction_new_models.py:480:12: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_extraction_new_models.py:481:12: E221 multiple spaces before operator
@@ -6756,25 +6916,24 @@
./tests/tfscreen/tfmodel/test_extraction_robustness.py:224:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_extraction_robustness.py:268:5: F841 local variable 'first_param' is assigned to but never used
./tests/tfscreen/tfmodel/test_extraction_robustness.py:270:5: F841 local variable 'param_name' is assigned to but never used
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:281:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:284:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:287:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:293:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:298:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:311:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:315:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:334:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:374:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:380:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:387:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_extraction_robustness.py:390:9: F841 local variable 'ds' is assigned to but never used
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:282:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:285:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:288:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:294:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:299:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:312:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:316:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:335:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:375:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:381:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:388:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_extraction_robustness.py:391:9: F841 local variable 'ds' is assigned to but never used
./tests/tfscreen/tfmodel/test_fit_model_helpers.py:16:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_fit_model_helpers.py:42:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_fit_model_helpers.py:87:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_fit_model_helpers.py:103:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_fit_model_helpers.py:113:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_fit_model_helpers.py:125:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_fit_model_helpers.py:162:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_fit_model_helpers.py:78:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_fit_model_helpers.py:88:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_fit_model_helpers.py:100:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_fit_model_helpers.py:135:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:5:1: F401 'numpyro.handlers' imported but unused
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:7:1: F401 'numpyro.infer.SVI' imported but unused
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:7:1: F401 'numpyro.infer.Trace_ELBO' imported but unused
@@ -6811,38 +6970,38 @@
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:110:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:118:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:126:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:129:5: F841 local variable 'num_titrants' is assigned to but never used
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:132:5: E306 expected 1 blank line before a nested definition, found 0
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:137:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:140:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:143:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:145:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:146:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:147:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:147:48: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:148:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:148:46: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:149:24: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:150:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:152:54: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:154:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:158:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:163:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:170:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:176:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:178:5: F841 local variable 'svi' is assigned to but never used
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:179:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:183:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:188:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:176:5: F841 local variable 'num_titrants' is assigned to but never used
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:179:5: E306 expected 1 blank line before a nested definition, found 0
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:184:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:187:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:190:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:192:59: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:193:28: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:193:52: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:194:28: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:194:50: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:195:28: E127 continuation line over-indented for visual indent
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:196:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:200:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:192:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:193:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:194:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:194:48: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:195:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:195:46: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:196:24: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:197:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:199:54: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:201:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:205:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:210:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:217:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:223:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:225:5: F841 local variable 'svi' is assigned to but never used
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:226:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:230:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:235:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:237:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:239:59: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:240:28: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:240:52: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:241:28: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:241:50: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:242:28: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:243:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_get_posteriors_mapping.py:247:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_h5_indexing.py:4:1: F401 'os' imported but unused
./tests/tfscreen/tfmodel/test_h5_indexing.py:9:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_h5_indexing.py:11:1: W293 blank line contains whitespace
@@ -6963,79 +7122,76 @@
./tests/tfscreen/tfmodel/test_model_orchestrator.py:103:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_model_orchestrator.py:109:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_model_orchestrator.py:113:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:120:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:129:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:146:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:149:51: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:155:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:159:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:241:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:264:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:270:128: E501 line too long (129 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:272:76: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:125:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:148:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:165:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_model_orchestrator.py:295:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:303:128: E501 line too long (184 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:304:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:314:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:319:48: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:319:50: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:319:52: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:319:54: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:319:56: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:319:58: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:320:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:328:112: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:328:128: E501 line too long (191 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:333:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:336:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:340:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:347:128: E501 line too long (184 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:366:128: E501 line too long (198 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:371:128: E501 line too long (167 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:419:79: E712 comparison to False should be 'if cond is False:' or 'if not cond:'
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:420:64: E712 comparison to True should be 'if cond is True:' or 'if cond:'
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:423:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:431:51: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:437:77: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:437:98: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:437:125: E231 missing whitespace after ':'
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:437:128: E501 line too long (131 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:438:128: E501 line too long (138 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:440:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:442:86: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:454:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:459:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:467:128: E501 line too long (129 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:475:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:486:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:489:128: E501 line too long (129 > 127 characters)
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:496:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:503:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:511:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:516:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:521:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:527:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:536:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:540:26: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:550:30: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:551:25: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:553:34: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:554:33: W291 trailing whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:561:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:566:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:43: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:45: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:47: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:49: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:51: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:53: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:568:55: E231 missing whitespace after ','
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:572:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:581:26: E261 at least two spaces before inline comment
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:582:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:585:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:588:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_model_orchestrator.py:765:67: W292 no newline at end of file
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:318:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:324:128: E501 line too long (129 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:326:76: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:349:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:357:128: E501 line too long (184 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:358:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:368:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:373:48: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:373:50: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:373:52: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:373:54: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:373:56: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:373:58: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:374:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:382:112: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:382:128: E501 line too long (191 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:387:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:390:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:394:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:401:128: E501 line too long (184 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:420:128: E501 line too long (198 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:425:128: E501 line too long (167 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:473:79: E712 comparison to False should be 'if cond is False:' or 'if not cond:'
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:474:64: E712 comparison to True should be 'if cond is True:' or 'if cond:'
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:477:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:485:51: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:491:77: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:491:98: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:491:125: E231 missing whitespace after ':'
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:491:128: E501 line too long (131 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:492:128: E501 line too long (138 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:494:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:496:86: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:508:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:513:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:521:128: E501 line too long (129 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:529:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:540:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:543:128: E501 line too long (129 > 127 characters)
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:550:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:557:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:565:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:570:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:575:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:581:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:590:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:594:26: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:604:30: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:605:25: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:607:34: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:608:33: W291 trailing whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:620:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:43: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:45: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:47: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:49: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:51: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:53: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:622:55: E231 missing whitespace after ','
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:626:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:635:26: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:636:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:639:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:642:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:899:9: F401 'inspect' imported but unused
+./tests/tfscreen/tfmodel/test_model_orchestrator.py:969:72: W292 no newline at end of file
./tests/tfscreen/tfmodel/test_needs_mut_wiring.py:30:16: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_needs_mut_wiring.py:30:31: E272 multiple spaces before keyword
./tests/tfscreen/tfmodel/test_needs_mut_wiring.py:31:16: E221 multiple spaces before operator
@@ -7057,7 +7213,6 @@
./tests/tfscreen/tfmodel/test_populate_dataclass.py:133:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_populate_dataclass.py:137:44: W292 no newline at end of file
./tests/tfscreen/tfmodel/test_posteriors.py:6:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_posteriors.py:10:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_posteriors.py:16:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_posteriors.py:20:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_posteriors.py:22:1: W293 blank line contains whitespace
@@ -7066,10 +7221,7 @@
./tests/tfscreen/tfmodel/test_posteriors.py:37:14: E261 at least two spaces before inline comment
./tests/tfscreen/tfmodel/test_posteriors.py:38:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_posteriors.py:45:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_posteriors.py:50:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_posteriors.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_posteriors.py:57:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_posteriors.py:60:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_posteriors.py:65:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_posteriors.py:72:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_posteriors.py:75:1: W293 blank line contains whitespace
@@ -7080,18 +7232,30 @@
./tests/tfscreen/tfmodel/test_posteriors.py:101:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_posteriors.py:108:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_posteriors.py:138:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:3:1: F401 'numpy as np' imported but unused
-./tests/tfscreen/tfmodel/test_prediction.py:7:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:33:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:42:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:54:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:64:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:68:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_prediction.py:71:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_prediction.py:75:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:85:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/tfmodel/test_prediction.py:89:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_prediction.py:91:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_prediction.py:4:1: F401 'jax.numpy as jnp' imported but unused
+./tests/tfscreen/tfmodel/test_prediction.py:6:1: F401 'unittest.mock.MagicMock' imported but unused
+./tests/tfscreen/tfmodel/test_prediction.py:6:1: F401 'unittest.mock.patch' imported but unused
+./tests/tfscreen/tfmodel/test_prediction.py:11:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_prediction.py:37:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_prediction.py:46:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_prediction.py:58:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_prediction.py:68:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_prediction.py:72:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_prediction.py:75:1: W293 blank line contains whitespace
+./tests/tfscreen/tfmodel/test_prediction.py:79:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_prediction.py:89:1: E302 expected 2 blank lines, found 1
+./tests/tfscreen/tfmodel/test_presplit.py:12:1: F401 'numpy as np' imported but unused
+./tests/tfscreen/tfmodel/test_presplit.py:15:1: F401 'numpyro' imported but unused
+./tests/tfscreen/tfmodel/test_presplit.py:24:1: F401 'tfscreen.tfmodel.data_class.DataClass' imported but unused
+./tests/tfscreen/tfmodel/test_presplit.py:25:1: F401 'tfscreen.tfmodel.tensors.populate_dataclass.populate_dataclass' imported but unused
+./tests/tfscreen/tfmodel/test_presplit.py:93:52: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:96:33: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:126:35: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:139:48: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:173:47: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:199:28: E128 continuation line under-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:217:47: E127 continuation line over-indented for visual indent
+./tests/tfscreen/tfmodel/test_presplit.py:229:28: E128 continuation line under-indented for visual indent
./tests/tfscreen/tfmodel/test_reference_math.py:29:11: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_reference_math.py:32:9: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_reference_math.py:75:14: E221 multiple spaces before operator
@@ -7291,13 +7455,12 @@
./tests/tfscreen/tfmodel/test_run_inference.py:903:25: E127 continuation line over-indented for visual indent
./tests/tfscreen/tfmodel/test_run_inference.py:903:26: E201 whitespace after '['
./tests/tfscreen/tfmodel/test_run_inference.py:904:26: E201 whitespace after '['
-./tests/tfscreen/tfmodel/test_sbc.py:6:1: F401 'tempfile' imported but unused
./tests/tfscreen/tfmodel/test_share_replicates.py:1:1: F401 'pytest' imported but unused
./tests/tfscreen/tfmodel/test_share_replicates.py:3:1: F401 'numpy as np' imported but unused
./tests/tfscreen/tfmodel/test_share_replicates.py:7:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/tfmodel/test_share_replicates.py:10:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_share_replicates.py:32:1: W293 blank line contains whitespace
-./tests/tfscreen/tfmodel/test_share_replicates.py:35:77: E261 at least two spaces before inline comment
+./tests/tfscreen/tfmodel/test_share_replicates.py:35:78: E261 at least two spaces before inline comment
./tests/tfscreen/tfmodel/test_share_replicates.py:40:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_share_replicates.py:42:1: W293 blank line contains whitespace
./tests/tfscreen/tfmodel/test_share_replicates.py:45:1: W293 blank line contains whitespace
@@ -7338,7 +7501,7 @@
./tests/tfscreen/tfmodel/test_struct_wiring.py:195:16: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_struct_wiring.py:197:12: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_struct_wiring.py:226:9: F401 'jax.numpy as jnp' imported but unused
-./tests/tfscreen/tfmodel/test_struct_wiring.py:233:61: E272 multiple spaces before keyword
+./tests/tfscreen/tfmodel/test_struct_wiring.py:233:71: E272 multiple spaces before keyword
./tests/tfscreen/tfmodel/test_struct_wiring.py:247:20: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_struct_wiring.py:249:16: E221 multiple spaces before operator
./tests/tfscreen/tfmodel/test_struct_wiring.py:283:5: F401 'os' imported but unused
@@ -7518,7 +7681,7 @@
./tests/tfscreen/util/io/test_read_dataframe.py:192:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/util/io/test_read_dataframe.py:200:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/util/io/test_read_dataframe.py:207:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/util/io/test_read_dataframe.py:215:40: W292 no newline at end of file
+./tests/tfscreen/util/io/test_read_dataframe.py:215:67: W292 no newline at end of file
./tests/tfscreen/util/numerical/test_array_search.py:78:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/util/numerical/test_array_search.py:143:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/util/numerical/test_array_search.py:167:1: E302 expected 2 blank lines, found 1
@@ -7594,54 +7757,53 @@
./tests/tfscreen/util/numerical/test_zero_truncated_poisson.py:81:1: W293 blank line contains whitespace
./tests/tfscreen/util/validation/test_check.py:8:1: E302 expected 2 blank lines, found 1
./tests/tfscreen/util/validation/test_check.py:26:1: E302 expected 2 blank lines, found 1
-./tests/tfscreen/util/validation/test_check.py:42:38: W292 no newline at end of file
-45 C901 'cat_fit' is too complex (19)
-16 E111 indentation is not a multiple of 4
+./tests/tfscreen/util/validation/test_check.py:75:70: W292 no newline at end of file
+50 C901 'cat_fit' is too complex (19)
+15 E111 indentation is not a multiple of 4
6 E114 indentation is not a multiple of 4 (comment)
6 E116 unexpected indentation (comment)
-12 E117 over-indented
+11 E117 over-indented
1 E122 continuation line missing indentation or outdented
2 E124 closing bracket does not match visual indentation
6 E125 continuation line with same indent as next logical line
-153 E127 continuation line over-indented for visual indent
-66 E128 continuation line under-indented for visual indent
+180 E127 continuation line over-indented for visual indent
+79 E128 continuation line under-indented for visual indent
3 E131 continuation line unaligned for hanging indent
43 E201 whitespace after '['
28 E202 whitespace before ']'
6 E203 whitespace before ','
5 E211 whitespace before '('
-1711 E221 multiple spaces before operator
+1832 E221 multiple spaces before operator
26 E222 multiple spaces after operator
8 E225 missing whitespace around operator
-1112 E231 missing whitespace after ','
-79 E251 unexpected spaces around keyword / parameter equals
+1115 E231 missing whitespace after ','
+103 E251 unexpected spaces around keyword / parameter equals
14 E252 missing whitespace around parameter equals
-226 E261 at least two spaces before inline comment
+229 E261 at least two spaces before inline comment
1 E262 inline comment should start with '# '
2 E265 block comment should start with '# '
2 E266 too many leading '#' for block comment
-150 E272 multiple spaces before keyword
+149 E272 multiple spaces before keyword
1 E301 expected 1 blank line, found 0
801 E302 expected 2 blank lines, found 1
-51 E303 too many blank lines (2)
-17 E305 expected 2 blank lines after class or function definition, found 1
+50 E303 too many blank lines (2)
+16 E305 expected 2 blank lines after class or function definition, found 1
29 E306 expected 1 blank line before a nested definition, found 0
1 E401 multiple imports on one line
-32 E402 module level import not at top of file
-71 E501 line too long (145 > 127 characters)
+37 E402 module level import not at top of file
+84 E501 line too long (145 > 127 characters)
46 E701 multiple statements on one line (colon)
-1 E702 multiple statements on one line (semicolon)
5 E712 comparison to True should be 'if cond is True:' or 'if cond:'
2 E713 test for membership should be 'not in'
1 E722 do not use bare 'except'
1 E731 do not assign a lambda expression, use a def
-12 E741 ambiguous variable name 'l'
-214 F401 'tfscreen.tfmodel.generative.components.theta.thermo.O2_C12_K5_U0_a.thermo.run_model' imported but unused
+20 E741 ambiguous variable name 'l'
+231 F401 'tfscreen.tfmodel.generative.components.theta.thermo.O2_C12_K5_U0_a.thermo.run_model' imported but unused
11 F541 f-string is missing placeholders
-3 F811 redefinition of unused 'np' from line 1
-40 F841 local variable 'e' is assigned to but never used
-693 W291 trailing whitespace
+4 F811 redefinition of unused 'np' from line 1
+41 F841 local variable 'e' is assigned to but never used
+669 W291 trailing whitespace
86 W292 no newline at end of file
-1721 W293 blank line contains whitespace
+1673 W293 blank line contains whitespace
29 W391 blank line at end of file
-7597
+7760
diff --git a/reports/junit/junit.xml b/reports/junit/junit.xml
index fd03d45a..f9e68b49 100644
--- a/reports/junit/junit.xml
+++ b/reports/junit/junit.xml
@@ -1 +1 @@
-
\ No newline at end of file
+
\ No newline at end of file
diff --git a/src/tfscreen/__version__.py b/src/tfscreen/__version__.py
index 663dd15d..ff3c18a2 100644
--- a/src/tfscreen/__version__.py
+++ b/src/tfscreen/__version__.py
@@ -5,5 +5,5 @@
Version string.
"""
-VERSION = (0, 4, 0)
+VERSION = (0, 4, 1)
__version__ = '.'.join(map(str, VERSION))
diff --git a/src/tfscreen/analysis/__init__.py b/src/tfscreen/analysis/__init__.py
index f63cb8a2..8175731d 100644
--- a/src/tfscreen/analysis/__init__.py
+++ b/src/tfscreen/analysis/__init__.py
@@ -7,3 +7,7 @@
extract_epistasis
)
+from .stats_test_suite import ( # noqa: F401
+ stats_test_suite
+)
+
diff --git a/src/tfscreen/analysis/cat_response/fit_response_cli.py b/src/tfscreen/analysis/cat_response/fit_response_cli.py
deleted file mode 100644
index 4f9541cc..00000000
--- a/src/tfscreen/analysis/cat_response/fit_response_cli.py
+++ /dev/null
@@ -1,186 +0,0 @@
-"""
-CLI for fitting cat_response models to theta-vs-titrant data per genotype.
-"""
-
-import argparse
-from concurrent.futures import ProcessPoolExecutor, as_completed
-
-import pandas as pd
-
-from tfscreen.analysis.cat_response.cat_fit import cat_fit
-from tfscreen.mle.curve_models import MODEL_LIBRARY
-
-DEFAULT_MODELS = ["flat", "repressor", "inducer", "hill_repressor", "hill_inducer"]
-
-
-def _fit_one(args):
- """Worker: run cat_fit for one (genotype, titrant_name) pair."""
- genotype, titrant_name, x, y, y_std, models_to_run = args
- flat_out, _ = cat_fit(x, y, y_std, models_to_run=models_to_run)
- flat_out["genotype"] = genotype
- flat_out["titrant_name"] = titrant_name
- return flat_out
-
-
-def fit_response(df,
- theta_col=None,
- sigma_col=None,
- models_to_run=None,
- workers=1):
- """
- Run cat_fit for every (genotype, titrant_name) group in df.
-
- Parameters
- ----------
- df : pd.DataFrame
- Must contain: genotype, titrant_name, titrant_conc, theta_col, and
- either sigma_col or both upper_std and lower_std.
- theta_col : str or None
- Column holding theta values passed to the fitter as y. If ``None``,
- auto-detected: ``median`` if present, then ``point_est``.
- sigma_col : str or None
- Column holding per-observation sigma. If None and df has
- upper_std / lower_std, sigma = (upper_std - lower_std) / 2.
- models_to_run : list of str or None
- Subset of MODEL_LIBRARY keys to fit. Defaults to DEFAULT_MODELS.
- workers : int
- Parallel worker processes.
-
- Returns
- -------
- pd.DataFrame
- One row per (genotype, titrant_name). Columns include best_model,
- status, AIC weights, and all parameter estimates/stds for every model,
- with '|' replaced by '_' in column names.
- """
- if models_to_run is None:
- models_to_run = DEFAULT_MODELS
-
- bad = [m for m in models_to_run if m not in MODEL_LIBRARY]
- if bad:
- raise ValueError(f"Unknown model(s): {bad}. Valid: {list(MODEL_LIBRARY)}")
-
- if theta_col is None:
- if "median" in df.columns:
- theta_col = "median"
- elif "point_est" in df.columns:
- theta_col = "point_est"
- else:
- raise ValueError(
- "No theta column found. Expected 'median' (posterior) or "
- "'point_est' (MAP). Pass theta_col explicitly to override."
- )
-
- if sigma_col is None:
- if "upper_std" in df.columns and "lower_std" in df.columns:
- df = df.copy()
- df["_sigma"] = (df["upper_std"] - df["lower_std"]) / 2
- sigma_col = "_sigma"
- else:
- raise ValueError(
- "No sigma_col specified and df lacks upper_std/lower_std columns. "
- "Provide --sigma_col or ensure upper_std and lower_std are present."
- )
-
- work_items = []
- for (genotype, titrant_name), group in df.groupby(["genotype", "titrant_name"],
- sort=False):
- x = group["titrant_conc"].to_numpy(dtype=float)
- y = group[theta_col].to_numpy(dtype=float)
- y_std = group[sigma_col].to_numpy(dtype=float)
- work_items.append((genotype, titrant_name, x, y, y_std, models_to_run))
-
- n_total = len(work_items)
- results = [None] * n_total
- idx_map = {id(item): i for i, item in enumerate(work_items)}
-
- print(f" Fitting {n_total} (genotype, titrant_name) pairs "
- f"with {workers} worker(s)...", flush=True)
-
- with ProcessPoolExecutor(max_workers=workers) as executor:
- futures = {executor.submit(_fit_one, item): idx_map[id(item)]
- for item in work_items}
- n_done = 0
- for future in as_completed(futures):
- results[futures[future]] = future.result()
- n_done += 1
- if n_done % 5000 == 0 or n_done == n_total:
- print(f" {n_done}/{n_total} fits complete", flush=True)
-
- out_df = pd.DataFrame(results)
- id_cols = ["genotype", "titrant_name"]
- other_cols = [c for c in out_df.columns if c not in id_cols]
- out_df = out_df[id_cols + other_cols].copy()
- out_df.columns = [c.replace("|", "_") for c in out_df.columns]
- return out_df
-
-
-def main():
- """CLI entry point for fitting cat_response models to theta data."""
-
- parser = argparse.ArgumentParser(
- prog="tfs-fit-response",
- description=fit_response.__doc__,
- formatter_class=argparse.RawTextHelpFormatter,
- )
- parser.add_argument(
- "input",
- type=str,
- help="Input CSV with columns: genotype, titrant_name, titrant_conc, "
- "and the theta / sigma columns.",
- )
- parser.add_argument(
- "--theta_col",
- type=str,
- default=None,
- help="Column name for theta values. Auto-detected if omitted "
- "('median' if present, else 'point_est').",
- )
- parser.add_argument(
- "--sigma_col",
- type=str,
- default=None,
- help="Column name for per-observation sigma. If omitted, computed "
- "as (upper_std - lower_std) / 2.",
- )
- parser.add_argument(
- "--models",
- type=str,
- nargs="+",
- default=None,
- metavar="MODEL",
- help="Models to fit (default: flat repressor inducer hill_repressor "
- "hill_inducer). Valid choices: " + " ".join(MODEL_LIBRARY),
- )
- parser.add_argument(
- "--workers",
- type=int,
- default=1,
- help="Number of parallel worker processes (default: 1).",
- )
- parser.add_argument(
- "--out",
- type=str,
- default="fit_response.csv",
- help="Output CSV path (default: fit_response.csv).",
- )
-
- args = parser.parse_args()
-
- print(f"Reading {args.input}...", flush=True)
- df = pd.read_csv(args.input)
-
- result_df = fit_response(
- df,
- theta_col=args.theta_col,
- sigma_col=args.sigma_col,
- models_to_run=args.models,
- workers=args.workers,
- )
-
- result_df.to_csv(args.out, index=False)
- print(f"Wrote {len(result_df)} rows to {args.out}", flush=True)
-
-
-if __name__ == "__main__":
- main()
diff --git a/tests/tfscreen/tfmodel/components/activity/__init__.py b/src/tfscreen/analysis/cat_response/scripts/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/activity/__init__.py
rename to src/tfscreen/analysis/cat_response/scripts/__init__.py
diff --git a/src/tfscreen/analysis/cat_response/cat_response_cli.py b/src/tfscreen/analysis/cat_response/scripts/cat_response_cli.py
similarity index 91%
rename from src/tfscreen/analysis/cat_response/cat_response_cli.py
rename to src/tfscreen/analysis/cat_response/scripts/cat_response_cli.py
index 8747e559..236f58fb 100644
--- a/src/tfscreen/analysis/cat_response/cat_response_cli.py
+++ b/src/tfscreen/analysis/cat_response/scripts/cat_response_cli.py
@@ -65,25 +65,25 @@ def cat_response(theta_file,
df = pd.read_csv(theta_file)
if theta_col is None:
- if "median" in df.columns:
- theta_col = "median"
+ if "q0.5" in df.columns:
+ theta_col = "q0.5"
elif "point_est" in df.columns:
theta_col = "point_est"
else:
raise ValueError(
- "No theta column found. Expected 'median' (posterior) or "
+ "No theta column found. Expected 'q0.5' (posterior median) or "
"'point_est' (MAP). Use --theta_col to specify a column explicitly."
)
if sigma_col is None:
- if "upper_std" in df.columns and "lower_std" in df.columns:
+ if "q0.841" in df.columns and "q0.159" in df.columns:
df = df.copy()
- df["_sigma"] = (df["upper_std"] - df["lower_std"]) / 2
+ df["_sigma"] = (df["q0.841"] - df["q0.159"]) / 2
sigma_col = "_sigma"
else:
raise ValueError(
- "No sigma_col specified and df lacks upper_std/lower_std columns. "
- "Provide --sigma_col or ensure upper_std and lower_std are present."
+ "No sigma_col specified and df lacks q0.841/q0.159 columns. "
+ "Provide --sigma_col or ensure q0.841 and q0.159 are present."
)
work_items = []
diff --git a/src/tfscreen/mle/stats_test_suite.py b/src/tfscreen/analysis/stats_test_suite.py
similarity index 100%
rename from src/tfscreen/mle/stats_test_suite.py
rename to src/tfscreen/analysis/stats_test_suite.py
diff --git a/src/tfscreen/genetics/library_manager.py b/src/tfscreen/genetics/library_manager.py
index 98b2fb17..e1a9c67c 100644
--- a/src/tfscreen/genetics/library_manager.py
+++ b/src/tfscreen/genetics/library_manager.py
@@ -250,15 +250,15 @@ def _parse_and_validate(self,
# -- Validate the sequence/library specification --
# Load wildtype seq
- wt_seq = str(run_config["wt_seq"]).strip()
+ wt_seq = "".join(str(run_config["wt_seq"]).split())
_check_char(wt_seq,"wt_seq",self.standard_bases)
-
+
# Load degenerate sites
- degen_sites = str(run_config["degen_sites"]).strip()
+ degen_sites = "".join(str(run_config["degen_sites"]).split())
_check_char(degen_sites,"degen_sites",self.degen_plus_dot)
-
+
# Load sub-libraries
- sub_libraries = str(run_config["sub_libraries"]).strip()
+ sub_libraries = "".join(str(run_config["sub_libraries"]).split())
_check_contiguous_lib_blocks(sub_libraries)
libraries_seen = set(list(sub_libraries)) - {"."}
@@ -420,7 +420,9 @@ def _prepare_indiv_lib_blocks(self, lib_to_get: str) -> Tuple[List[List[str]], L
end_idx = indexes[-1] + 1
# Get the degenerate sites and wildtype sequences for this sub-library
- lib_seq = "".join(self.degen_sites[start_idx:end_idx])
+ lib_seq = "".join(d if d != '.' else w
+ for d, w in zip(self.degen_sites[start_idx:end_idx],
+ self.wt_seq[start_idx:end_idx]))
wt_seq = "".join(self.wt_seq[start_idx:end_idx])
# Extract the region of the library that encodes the degenerate library,
@@ -734,11 +736,14 @@ def _get_spiked_seqs(self,list_of_seqs):
spiked_seqs = []
for seq in list_of_seqs:
- _check_char(seq,"spiked seq",self.standard_bases)
+ seq = "".join(str(seq).split())
+ _check_char(seq,"spiked seq",self.standard_plus_dot)
if len(seq) != len(self.wt_seq):
raise ValueError(
f"spiked seq {seq} is not the same length as the library"
)
+ seq = "".join(w if s == "." else s
+ for s, w in zip(seq, self.wt_seq))
spiked_seqs.append(seq)
spiked_aa = self._convert_to_aa(spiked_seqs)
diff --git a/src/tfscreen/mle/__init__.py b/src/tfscreen/mle/__init__.py
index 47e0cc43..ede64f9f 100644
--- a/src/tfscreen/mle/__init__.py
+++ b/src/tfscreen/mle/__init__.py
@@ -9,9 +9,6 @@
predict_with_error
)
-from .stats_test_suite import ( # noqa: F401
- stats_test_suite
-)
from .fit_manager import ( # noqa: F401
FitManager
diff --git a/src/tfscreen/plot/__init__.py b/src/tfscreen/plot/__init__.py
index 1a657d60..8a63e847 100644
--- a/src/tfscreen/plot/__init__.py
+++ b/src/tfscreen/plot/__init__.py
@@ -42,3 +42,7 @@
from .heatmap import ( # noqa: F401
heatmap
)
+
+from .geno_trajectory import ( # noqa: F401
+ plot_geno_trajectory,
+)
diff --git a/src/tfscreen/plot/geno_trajectory.py b/src/tfscreen/plot/geno_trajectory.py
new file mode 100644
index 00000000..9662be97
--- /dev/null
+++ b/src/tfscreen/plot/geno_trajectory.py
@@ -0,0 +1,400 @@
+"""
+Functions for plotting per-genotype growth trajectory predictions.
+
+plot_geno_trajectory
+ Pure plotting function. Accepts a prediction DataFrame and returns a
+ ``matplotlib.figure.Figure``.
+
+predict_and_plot_geno_trajectory
+ Composite entry point. Calls ``predict()`` with a fine t_sel grid, merges
+ in observed data and a pre-selection anchor from the ln_cfu0 site, then
+ delegates to ``plot_geno_trajectory``.
+"""
+
+import numpy as np
+import pandas as pd
+
+_T_FINE = 50
+
+_CONDITION_COLS = [
+ "condition_pre", "condition_sel", "titrant_name", "titrant_conc"
+]
+_KEEP_COLS = [
+ "replicate", "condition_pre", "condition_sel",
+ "titrant_name", "titrant_conc", "genotype",
+ "t_sel", "ln_cfu", "ln_cfu_std", "q0.05", "q0.5", "q0.95",
+]
+
+
+def plot_geno_trajectory(
+ pred_df,
+ figsize=None,
+ colors=None,
+ markers=None,
+):
+ """
+ Plot per-genotype growth trajectories from a prediction DataFrame.
+
+ One subplot is generated per unique
+ ``(condition_pre, condition_sel, titrant_name, titrant_conc)`` combination.
+ Every ``(genotype, replicate)`` pair present in ``pred_df`` is drawn as a
+ separate series on each subplot. Genotypes are distinguished by color;
+ replicates are distinguished by marker and linestyle.
+
+ Parameters
+ ----------
+ pred_df : pd.DataFrame
+ Prediction DataFrame. Required columns:
+
+ * ``replicate``, ``condition_pre``, ``condition_sel``,
+ ``titrant_name``, ``titrant_conc``, ``genotype`` — grouping keys.
+ * ``t_sel`` — time coordinate (may include negative values for the
+ pre-selection phase).
+ * ``ln_cfu`` — observed ln(CFU); ``NaN`` where no observation exists.
+ * ``ln_cfu_std`` — observed std; ``NaN`` where no observation exists.
+ * ``q0.5`` — predicted median trajectory; ``NaN`` where unavailable.
+
+ Optional columns (both must be present together to draw a credible
+ interval):
+
+ * ``q0.05``, ``q0.95`` — 5th/95th percentile of the posterior predictive.
+
+ figsize : tuple of (float, float), optional
+ Figure ``(width, height)`` in inches. Defaults to
+ ``(5 * n_cols, 4 * n_rows)``.
+ colors : list of str, optional
+ Color cycle for genotypes. Defaults to ``DEFAULT_COLORS``.
+ markers : list of str, optional
+ Marker cycle for replicates. Defaults to ``DEFAULT_MARKERS``.
+
+ Returns
+ -------
+ matplotlib.figure.Figure
+ """
+ from matplotlib import pyplot as plt
+ from tfscreen.plot.default_styles import DEFAULT_COLORS, DEFAULT_MARKERS
+
+ if colors is None:
+ colors = DEFAULT_COLORS[:]
+ if markers is None:
+ markers = DEFAULT_MARKERS[:]
+
+ conditions = (
+ pred_df[_CONDITION_COLS]
+ .drop_duplicates()
+ .sort_values(_CONDITION_COLS)
+ .reset_index(drop=True)
+ )
+ genotypes = sorted(pred_df["genotype"].unique().tolist(), key=str)
+ replicates = sorted(pred_df["replicate"].unique().tolist(), key=str)
+
+ n_combos = len(conditions)
+ n_cols = min(3, n_combos)
+ n_rows = (n_combos + n_cols - 1) // n_cols
+
+ if figsize is None:
+ figsize = (5 * n_cols, 4 * n_rows)
+
+ fig, axes = plt.subplots(
+ n_rows, n_cols, figsize=figsize, squeeze=False, sharey=True
+ )
+
+ # One color per (genotype, replicate) pair
+ series_pairs = [
+ (g, r) for g in genotypes for r in replicates
+ ]
+ pair_color = {
+ pair: colors[i % len(colors)] for i, pair in enumerate(series_pairs)
+ }
+ rep_marker = {r: markers[i % len(markers)] for i, r in enumerate(replicates)}
+
+ has_ci = {"q0.05", "q0.95"} <= set(pred_df.columns)
+
+ for combo_i, cond_row in conditions.iterrows():
+ ax = axes[combo_i // n_cols][combo_i % n_cols]
+
+ cp = cond_row["condition_pre"]
+ cs = cond_row["condition_sel"]
+ tn = cond_row["titrant_name"]
+ tc = float(cond_row["titrant_conc"])
+
+ cond_df = pred_df[
+ (pred_df["condition_pre"] == cp)
+ & (pred_df["condition_sel"] == cs)
+ & (pred_df["titrant_name"] == tn)
+ & (pred_df["titrant_conc"] == tc)
+ ].copy()
+
+ ax.set_title(f"{cp} → {cs}\n{tn} = {tc:.3g}", fontsize=9)
+ ax.set_xlabel("Time")
+ ax.set_ylabel("ln(CFU)")
+ ax.axvline(0.0, color="0.6", lw=0.8, ls="--")
+ ax.spines["top"].set_visible(False)
+ ax.spines["right"].set_visible(False)
+
+ for geno in genotypes:
+ geno_df = cond_df[cond_df["genotype"] == geno]
+ if geno_df.empty:
+ continue
+
+ for rep in replicates:
+ rep_df = (
+ geno_df[geno_df["replicate"] == rep]
+ .sort_values("t_sel")
+ )
+ if rep_df.empty:
+ continue
+
+ color = pair_color[(geno, rep)]
+ label = f"{geno} ({rep})"
+ mk = rep_marker[rep]
+ ls = "-"
+
+ # Observed data points
+ obs = rep_df.dropna(subset=["ln_cfu"])
+ if not obs.empty:
+ ax.errorbar(
+ obs["t_sel"],
+ obs["ln_cfu"],
+ yerr=obs["ln_cfu_std"],
+ fmt=mk,
+ color=color,
+ ms=5,
+ lw=1,
+ capsize=3,
+ zorder=3,
+ )
+
+ # Predicted median line
+ line_df = rep_df.dropna(subset=["q0.5"])
+ if not line_df.empty:
+ ax.plot(
+ line_df["t_sel"],
+ line_df["q0.5"],
+ ls,
+ color=color,
+ lw=1.8,
+ label=label,
+ zorder=4,
+ )
+
+ # 90 % credible interval band
+ if has_ci:
+ ci_df = rep_df.dropna(subset=["q0.05", "q0.95"])
+ if not ci_df.empty:
+ ax.fill_between(
+ ci_df["t_sel"],
+ ci_df["q0.05"],
+ ci_df["q0.95"],
+ color=color,
+ alpha=0.2,
+ zorder=2,
+ )
+
+ ax.legend(fontsize=8, loc="best")
+
+ for extra_i in range(n_combos, n_rows * n_cols):
+ axes[extra_i // n_cols][extra_i % n_cols].set_visible(False)
+
+ fig.tight_layout()
+ return fig
+
+
+def predict_geno_trajectory_df(
+ orchestrator,
+ param_posteriors,
+ genotypes=None,
+ titrant_names=None,
+ t_fine=_T_FINE,
+ num_marginal_samples=200,
+):
+ """
+ Run forward predictions and return the merged trajectory DataFrame.
+
+ Calls :func:`~tfscreen.tfmodel.analysis.prediction.predict` with a fine
+ ``t_sel`` grid, merges in observed data and a pre-selection anchor from
+ the ``ln_cfu0`` site, and returns a single DataFrame ready for
+ :func:`plot_geno_trajectory`. Only ``(condition_pre, condition_sel,
+ titrant_name, titrant_conc)`` combinations present in
+ ``orchestrator.growth_df`` are included.
+
+ Parameters
+ ----------
+ orchestrator : ModelOrchestrator
+ Fitted model orchestrator.
+ param_posteriors : dict or str
+ MAP parameter dict (keys ending in ``_auto_loc``), path to a
+ ``_params.npz`` file, or path to a ``_posterior.h5`` file.
+ genotypes : list of str, optional
+ Subset of genotypes to include. If ``None``, uses all genotypes in
+ the orchestrator.
+ titrant_names : list of str, optional
+ Subset of titrant names to include. If ``None``, uses all.
+ t_fine : int, optional
+ Number of equally-spaced selection-phase time points for the fine grid.
+ Defaults to ``_T_FINE`` (50).
+ num_marginal_samples : int, optional
+ Number of posterior samples to run through the model when computing
+ quantiles. Passed directly to :func:`~.prediction.predict`.
+ Defaults to 200.
+
+ Returns
+ -------
+ pd.DataFrame
+ Merged prediction DataFrame with columns matching ``_KEEP_COLS``.
+ """
+ from tfscreen.tfmodel.analysis.prediction import predict
+
+ if isinstance(param_posteriors, str) and param_posteriors.endswith(".npz"):
+ param_posteriors = dict(np.load(param_posteriors))
+
+ # Fine selection-phase time grid: 0 … max observed t_sel
+ gd = orchestrator.data.growth
+ good_mask = np.asarray(gd.good_mask)
+ t_sel_tensor = np.asarray(gd.t_sel)
+ max_t_sel = float(np.nanmax(t_sel_tensor[good_mask]))
+ t_sel_grid = np.linspace(0.0, max_t_sel, t_fine).tolist()
+
+ all_dfs = predict(
+ orchestrator,
+ param_posteriors,
+ predict_sites=["growth_pred", "ln_cfu0"],
+ q_to_get=[0.05, 0.5, 0.95],
+ num_samples=None,
+ num_marginal_samples=num_marginal_samples,
+ t_sel=t_sel_grid,
+ genotypes=genotypes,
+ )
+ fine_df = all_dfs["growth_pred"]
+ ln_cfu0_raw = all_dfs["ln_cfu0"]
+
+ # Build anchor rows at t_sel = -t_pre using ln_cfu0 quantiles.
+ # ln_cfu0 is indexed by (replicate, condition_pre, genotype); expand across
+ # all valid (condition_sel, titrant_name, titrant_conc) for each condition_pre.
+ t_pre_df = (
+ orchestrator.growth_df[["replicate", "condition_pre", "t_pre"]]
+ .drop_duplicates(subset=["replicate", "condition_pre"])
+ )
+ ln_cfu0_vals = (
+ ln_cfu0_raw[["replicate", "condition_pre", "genotype", "q0.05", "q0.5", "q0.95"]]
+ .drop_duplicates(subset=["replicate", "condition_pre", "genotype"])
+ )
+ valid_combos = orchestrator.growth_df[_CONDITION_COLS].drop_duplicates()
+ anchor_df = (
+ ln_cfu0_vals
+ .merge(valid_combos, on="condition_pre", how="inner")
+ .merge(t_pre_df, on=["replicate", "condition_pre"], how="left")
+ )
+ anchor_df["t_sel"] = -anchor_df["t_pre"]
+ anchor_df["ln_cfu"] = np.nan
+ anchor_df["ln_cfu_std"] = np.nan
+
+ # Overlay observed ln_cfu0 measurements when presplit_df is available.
+ presplit_df = getattr(orchestrator, "presplit_df", None)
+ if presplit_df is not None:
+ ps = presplit_df[
+ ["replicate", "condition_pre", "genotype", "ln_cfu", "ln_cfu_std"]
+ ].rename(columns={"ln_cfu": "_ps_ln", "ln_cfu_std": "_ps_ln_std"})
+ anchor_df = anchor_df.merge(
+ ps, on=["replicate", "condition_pre", "genotype"], how="left"
+ )
+ anchor_df["ln_cfu"] = anchor_df["_ps_ln"]
+ anchor_df["ln_cfu_std"] = anchor_df["_ps_ln_std"]
+ anchor_df = anchor_df.drop(columns=["_ps_ln", "_ps_ln_std"])
+
+ # Observed data from the orchestrator (no model predictions)
+ obs_cols = [
+ "replicate", "condition_pre", "condition_sel",
+ "titrant_name", "titrant_conc", "genotype",
+ "t_sel", "ln_cfu", "ln_cfu_std",
+ ]
+ obs_df = orchestrator.growth_df[obs_cols].copy()
+ if genotypes is not None:
+ obs_df = obs_df[obs_df["genotype"].isin(list(genotypes))]
+ for col in ("q0.05", "q0.5", "q0.95"):
+ obs_df[col] = np.nan
+
+ # Titrant-name filter (predict() has no titrant_name argument)
+ if titrant_names is not None:
+ tn_set = {str(t) for t in titrant_names}
+ fine_df = fine_df[fine_df["titrant_name"].isin(tn_set)]
+ obs_df = obs_df[obs_df["titrant_name"].isin(tn_set)]
+ anchor_df = anchor_df[anchor_df["titrant_name"].isin(tn_set)]
+
+ pred_df = pd.concat(
+ [fine_df[_KEEP_COLS], obs_df[_KEEP_COLS], anchor_df[_KEEP_COLS]],
+ ignore_index=True,
+ )
+
+ # Restrict to condition combos that actually exist in the training data.
+ # copy_orchestrator() produces a Cartesian product of categorical columns
+ # that can include invalid (condition_pre, condition_sel) pairings; this
+ # filter removes them.
+ valid_combos = (
+ orchestrator.growth_df[_CONDITION_COLS]
+ .drop_duplicates()
+ )
+ pred_df = pred_df.merge(valid_combos, on=_CONDITION_COLS, how="inner")
+ return pred_df
+
+
+def predict_and_plot_geno_trajectory(
+ orchestrator,
+ param_posteriors,
+ genotypes=None,
+ titrant_names=None,
+ t_fine=_T_FINE,
+ num_marginal_samples=200,
+ figsize=None,
+ colors=None,
+ markers=None,
+):
+ """
+ Run forward predictions and plot per-genotype growth trajectories.
+
+ Calls :func:`predict_geno_trajectory_df` to build the merged prediction
+ DataFrame, then delegates to :func:`plot_geno_trajectory`. Only
+ ``(condition_pre, condition_sel, titrant_name, titrant_conc)`` combinations
+ present in ``orchestrator.growth_df`` are included in the output.
+
+ Parameters
+ ----------
+ orchestrator : ModelOrchestrator
+ Fitted model orchestrator.
+ param_posteriors : dict or str
+ MAP parameter dict (keys ending in ``_auto_loc``), path to a
+ ``_params.npz`` file, or path to a ``_posterior.h5`` file.
+ genotypes : list of str, optional
+ Subset of genotypes to include. If ``None``, uses all genotypes in
+ the orchestrator.
+ titrant_names : list of str, optional
+ Subset of titrant names to include. If ``None``, uses all.
+ t_fine : int, optional
+ Number of equally-spaced selection-phase time points for the fine grid.
+ Defaults to ``_T_FINE`` (50).
+ num_marginal_samples : int, optional
+ Number of posterior samples to run through the model when computing
+ quantiles. Passed directly to :func:`~.prediction.predict`.
+ Defaults to 200.
+ figsize : tuple of (float, float), optional
+ Passed to :func:`plot_geno_trajectory`.
+ colors : list of str, optional
+ Passed to :func:`plot_geno_trajectory`.
+ markers : list of str, optional
+ Passed to :func:`plot_geno_trajectory`.
+
+ Returns
+ -------
+ matplotlib.figure.Figure
+ """
+ pred_df = predict_geno_trajectory_df(
+ orchestrator,
+ param_posteriors,
+ genotypes=genotypes,
+ titrant_names=titrant_names,
+ t_fine=t_fine,
+ num_marginal_samples=num_marginal_samples,
+ )
+ return plot_geno_trajectory(
+ pred_df, figsize=figsize, colors=colors, markers=markers
+ )
diff --git a/src/tfscreen/plot/plot_theta_fits.py b/src/tfscreen/plot/plot_theta_fits.py
index 1dcff60c..73e7d93a 100644
--- a/src/tfscreen/plot/plot_theta_fits.py
+++ b/src/tfscreen/plot/plot_theta_fits.py
@@ -115,32 +115,33 @@ def plot_theta_fits(df,
color=colors[counter],zorder=0)
# Central estimate (posterior median or MAP point estimate)
- center_col = "median" if "median" in g_df.columns else "point_est"
+ center_col = "q0.5" if "q0.5" in g_df.columns else "point_est"
ax.plot(g_df["titrant_conc"],
g_df[center_col],
lw=2,color=colors[counter])
- # Standard error
- if set(["lower_std","upper_std"]) <= set(g_df.columns):
-
+ # ±1 SD credible band (q0.159 / q0.841)
+ if set(["q0.159","q0.841"]) <= set(g_df.columns):
+
ax.fill_between(x=g_df["titrant_conc"],
- y1=g_df["lower_std"],
- y2=g_df["upper_std"],
+ y1=g_df["q0.159"],
+ y2=g_df["q0.841"],
color=colors[counter],
alpha=0.7,zorder=-10)
- # 95% CI
- if set(["lower_95","upper_95"]) <= set(g_df.columns):
-
+ # 95% CI (q0.025 / q0.975)
+ if set(["q0.025","q0.975"]) <= set(g_df.columns):
+
ax.fill_between(x=g_df["titrant_conc"],
- y1=g_df["lower_95"],
- y2=g_df["upper_95"],
+ y1=g_df["q0.025"],
+ y2=g_df["q0.975"],
color=colors[counter],
alpha=0.4,zorder=-20)
counter += 1
ax.set_xscale("log")
ax.set_xlabel("titrant conc (mM)")
+ ax.set_ylim(-0.05,1.05)
ax.set_ylabel("$\\theta$")
ax.legend()
ax.spines['top'].set_visible(False)
diff --git a/src/tfscreen/simulate/binding_params.py b/src/tfscreen/simulate/binding_params.py
new file mode 100644
index 00000000..201c8b55
--- /dev/null
+++ b/src/tfscreen/simulate/binding_params.py
@@ -0,0 +1,481 @@
+"""
+Utilities for injecting measured Hill binding parameters into the simulation.
+
+When real binding data (theta curves from fluorescence, SPR, ITC, etc.) are
+available for specific genotypes, supply a CSV of per-genotype Hill parameters
+via the ``binding_data.genotype_params_file`` config key. This module reads
+those parameters and computes theta-value overrides for injection into the
+simulation via the ``theta_gc_override`` mechanism in ``thermo_to_growth``.
+
+Supported theta components: ``hill_geno``, ``hill_mut``.
+
+CSV format (``genotype_params_file``)
+--------------------------------------
+Columns: ``genotype``, plus any subset of
+ ``theta_low``, ``theta_high``, ``log_hill_K``, ``hill_n``.
+
+One row per genotype. ``NaN`` in any parameter column → the WT reference
+value for that parameter is used (from ``theta_sim_priors`` in the config,
+or overridden by the ``wt`` row in the CSV itself).
+
+Example::
+
+ genotype,theta_low,theta_high,log_hill_K,hill_n
+ wt,0.99,0.01,-4.1,2.0
+ A47V,0.97,0.03,-3.8,1.8
+ K84L,,,−3.5,
+
+``hill_mut`` note
+-----------------
+Single-mutant rows are back-converted to per-mutation deltas (relative to the
+WT reference) and injected into a full ``hill_mut``-style simulation. Deltas
+for mutations absent from the CSV are drawn from the ``SimPriors`` as usual.
+Pairwise epistasis is sampled from the regularized horseshoe prior and applied
+consistently to the whole library. The function returns theta for every
+library genotype — including doubles assembled from measured singles — so
+downstream growth rates are fully consistent with the measured binding data.
+"""
+
+import warnings
+import numpy as np
+import pandas as pd
+from pathlib import Path
+from typing import Union
+
+_EPS = 1e-7
+_THETA_CLIP = 1e-4 # hard bounds for theta_low / theta_high from CSV
+_ZERO_CONC_SENTINEL = 1e-20
+
+HILL_PARAM_COLS = frozenset({"theta_low", "theta_high", "log_hill_K", "hill_n"})
+SUPPORTED_COMPONENTS = frozenset({"hill_geno", "hill_mut"})
+
+
+# ---------------------------------------------------------------------------
+# CSV I/O
+# ---------------------------------------------------------------------------
+
+def read_binding_genotype_params(csv_path: Union[str, Path]) -> dict:
+ """
+ Read a per-genotype Hill parameter CSV.
+
+ Parameters
+ ----------
+ csv_path : str or Path
+
+ Returns
+ -------
+ dict[str, dict[str, float]]
+ ``{genotype: {param: value}}``. ``NaN`` values are preserved as
+ ``float("nan")``; the caller fills them from the WT reference.
+
+ Raises
+ ------
+ ValueError
+ If the CSV lacks a ``genotype`` column, has no recognised parameter
+ columns, or contains unrecognised column names.
+ """
+ df = pd.read_csv(csv_path)
+
+ if "genotype" not in df.columns:
+ raise ValueError(
+ f"'{csv_path}' must have a 'genotype' column."
+ )
+
+ param_cols_present = [c for c in df.columns if c in HILL_PARAM_COLS]
+ if not param_cols_present:
+ raise ValueError(
+ f"'{csv_path}' must have at least one parameter column: "
+ f"{sorted(HILL_PARAM_COLS)}"
+ )
+
+ unknown = set(df.columns) - {"genotype"} - HILL_PARAM_COLS
+ if unknown:
+ raise ValueError(
+ f"Unrecognised columns in '{csv_path}': {sorted(unknown)}. "
+ f"Expected: genotype + subset of {sorted(HILL_PARAM_COLS)}"
+ )
+
+ result = {}
+ for _, row in df.iterrows():
+ g = str(row["genotype"])
+ params = {}
+ for col in param_cols_present:
+ v = row[col]
+ if pd.isna(v):
+ params[col] = float("nan")
+ else:
+ v = float(v)
+ if col in ("theta_low", "theta_high"):
+ clipped = float(np.clip(v, _THETA_CLIP, 1.0 - _THETA_CLIP))
+ if clipped != v:
+ warnings.warn(
+ f"'{csv_path}': {col} for genotype '{g}' "
+ f"({v:.6g}) is outside [{_THETA_CLIP}, "
+ f"{1.0 - _THETA_CLIP}] and has been clipped to "
+ f"{clipped:.6g}. theta values must lie strictly "
+ f"between 0 and 1.",
+ UserWarning,
+ stacklevel=2,
+ )
+ v = clipped
+ params[col] = v
+ result[g] = params
+
+ return result
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+def _logit(x):
+ xc = np.clip(np.asarray(x, dtype=float), _EPS, 1.0 - _EPS)
+ return np.log(xc / (1.0 - xc))
+
+
+def _to_log_conc(concs):
+ c = np.asarray(concs, dtype=float)
+ return np.log(np.where(c == 0.0, _ZERO_CONC_SENTINEL, c))
+
+
+def _hill_theta(theta_low, theta_high, log_hill_K, hill_n, log_conc):
+ """Evaluate Hill equation at log concentrations. Returns shape (len(log_conc),)."""
+ lc = np.asarray(log_conc, dtype=float)
+ occupancy = 1.0 / (1.0 + np.exp(-hill_n * (lc - log_hill_K)))
+ return float(theta_low) + (float(theta_high) - float(theta_low)) * occupancy
+
+
+def _fill_params_from_wt(params: dict, wt_params: dict) -> dict:
+ """Return a copy of *params* with NaN entries replaced from *wt_params*."""
+ filled = dict(params)
+ for k in list(HILL_PARAM_COLS):
+ v = filled.get(k)
+ if v is None or (isinstance(v, float) and np.isnan(v)):
+ if k in wt_params:
+ filled[k] = wt_params[k]
+ return filled
+
+
+def _wt_params_from_sim_priors(sim_priors) -> dict:
+ """
+ Extract WT reference in the CSV parameter space from a SimPriors object.
+
+ Works with both ``hill_geno.SimPriors`` and ``hill_mut.SimPriors``;
+ both expose ``wt_theta_low``, ``wt_theta_high``, ``wt_log_K``, ``wt_hill_n``.
+ """
+ return {
+ "theta_low": float(sim_priors.wt_theta_low),
+ "theta_high": float(sim_priors.wt_theta_high),
+ "log_hill_K": float(sim_priors.wt_log_K),
+ "hill_n": float(sim_priors.wt_hill_n),
+ }
+
+
+# ---------------------------------------------------------------------------
+# hill_geno override
+# ---------------------------------------------------------------------------
+
+def build_theta_gc_override_hill_geno(
+ params_dict: dict,
+ log_conc_growth: np.ndarray,
+ wt_params: dict,
+) -> dict:
+ """
+ Compute theta at growth concentrations for measured genotypes (``hill_geno``).
+
+ Each entry in *params_dict* contributes one key to the returned override
+ dict. ``NaN`` parameter values are filled from *wt_params*.
+
+ Parameters
+ ----------
+ params_dict : dict[str, dict[str, float]]
+ Output of :func:`read_binding_genotype_params`.
+ log_conc_growth : np.ndarray, shape (C,)
+ Log concentrations at growth conditions (from ``sim_data.log_titrant_conc``).
+ wt_params : dict
+ WT reference: ``{theta_low, theta_high, log_hill_K, hill_n}``.
+ From :func:`_wt_params_from_sim_priors`.
+
+ Returns
+ -------
+ theta_dict : dict[str, np.ndarray]
+ ``{genotype: theta_array}`` where theta_array has shape (C,).
+ effective_params : dict[str, dict[str, float]]
+ ``{genotype: {theta_low, theta_high, log_hill_K, hill_n}}`` — the
+ actual Hill parameters used for each genotype (NaN fields filled from
+ *wt_params*).
+ """
+ theta_dict = {}
+ effective_params = {}
+ for g, params in params_dict.items():
+ filled = _fill_params_from_wt(params, wt_params)
+ theta_dict[g] = _hill_theta(
+ filled["theta_low"], filled["theta_high"],
+ filled["log_hill_K"], filled["hill_n"],
+ log_conc_growth,
+ )
+ effective_params[g] = {k: filled[k] for k in HILL_PARAM_COLS}
+ return theta_dict, effective_params
+
+
+# ---------------------------------------------------------------------------
+# hill_mut override
+# ---------------------------------------------------------------------------
+
+def build_theta_gc_override_hill_mut(
+ params_dict: dict,
+ library_genotypes: list,
+ sim_data,
+ sim_priors,
+ log_conc: np.ndarray,
+ rng: np.random.Generator,
+) -> dict:
+ """
+ Run a ``hill_mut``-style simulation with per-mutation deltas pinned for
+ measured genotypes and return theta for every library genotype.
+
+ Single-mutation entries in *params_dict* are back-converted to per-mutation
+ deltas relative to the WT reference (CSV ``wt`` row wins; otherwise
+ *sim_priors* WT values are used). Mutations absent from *params_dict* draw
+ deltas from the prior as usual. Pairwise epistasis is sampled from the
+ regularized horseshoe prior for all pairs consistently. Multi-mutation
+ entries in *params_dict* override the assembled theta value directly.
+
+ Parameters
+ ----------
+ params_dict : dict[str, dict[str, float]]
+ library_genotypes : list[str]
+ Genotype strings in the same order as *sim_data* was built from.
+ sim_data : SimData
+ Must have ``mut_nnz_mut_idx``, ``mut_nnz_geno_idx``,
+ ``pair_nnz_pair_idx``, ``pair_nnz_geno_idx``, ``num_pair``.
+ sim_priors : hill_mut.SimPriors
+ Used as fallback WT reference and prior widths for unmeasured mutations.
+ log_conc : np.ndarray, shape (C,)
+ Log concentrations to evaluate theta at.
+ rng : np.random.Generator
+
+ Returns
+ -------
+ theta_dict : dict[str, np.ndarray]
+ ``{genotype: theta_array}`` for every genotype in *library_genotypes*.
+ Caller should add all entries to ``theta_gc_override``.
+ effective_params : dict[str, dict[str, float]]
+ ``{genotype: {theta_low, theta_high, log_hill_K, hill_n}}`` — the
+ effective Hill parameters used for each genotype in the simulation.
+ """
+ import jax.numpy as jnp
+ from tfscreen.genetics.build_mut_geno_matrix import (
+ build_mut_geno_matrix,
+ apply_mut_matrix,
+ apply_pair_matrix,
+ )
+
+ log_conc = np.asarray(log_conc, dtype=float)
+ G = len(library_genotypes)
+
+ # Get mutation labels in the same order as sim_data was built
+ mut_labels, pair_labels, _, _, _ = build_mut_geno_matrix(library_genotypes)
+ M = len(mut_labels)
+ P = len(pair_labels)
+
+ # WT reference in logit/log space ---
+ # If 'wt' row is in params_dict, it overrides sim_priors WT values.
+ wt_params = _wt_params_from_sim_priors(sim_priors)
+ if "wt" in params_dict:
+ wt_from_csv = _fill_params_from_wt(params_dict["wt"], wt_params)
+ wt_params = wt_from_csv # CSV WT wins
+
+ wt_logit_low = float(_logit(wt_params["theta_low"]))
+ wt_logit_high = float(_logit(wt_params["theta_high"]))
+ wt_logit_delta = wt_logit_high - wt_logit_low
+ wt_log_K = float(wt_params["log_hill_K"])
+ wt_log_n = float(np.log(wt_params["hill_n"]))
+
+ # Build per-mutation delta arrays (NaN = needs to be sampled from prior)
+ d_logit_low = np.full(M, np.nan)
+ d_logit_delta = np.full(M, np.nan)
+ d_log_K = np.full(M, np.nan)
+ d_log_n = np.full(M, np.nan)
+
+ def _parse_muts(g):
+ return [p for p in g.split("/") if p and p.lower() != "wt"]
+
+ mut_label_to_idx = {m: i for i, m in enumerate(mut_labels)}
+
+ for g, params in params_dict.items():
+ muts = _parse_muts(g)
+ if len(muts) != 1:
+ continue # skip wt and multi-mutants for delta back-calculation
+ m_label = muts[0]
+ if m_label not in mut_label_to_idx:
+ continue # this mutation is not in the library
+
+ m_idx = mut_label_to_idx[m_label]
+ filled = _fill_params_from_wt(params, wt_params)
+
+ g_logit_low = float(_logit(filled["theta_low"]))
+ g_logit_high = float(_logit(filled["theta_high"]))
+ g_logit_delta = g_logit_high - g_logit_low
+
+ d_logit_low[m_idx] = g_logit_low - wt_logit_low
+ d_logit_delta[m_idx] = g_logit_delta - wt_logit_delta
+ d_log_K[m_idx] = float(filled["log_hill_K"]) - wt_log_K
+ d_log_n[m_idx] = float(np.log(filled["hill_n"])) - wt_log_n
+
+ # Sample deltas for unmeasured mutations from prior
+ unmeasured = np.isnan(d_logit_low)
+ n_un = int(unmeasured.sum())
+ if n_un > 0:
+ d_logit_low[unmeasured] = rng.normal(0.0, sim_priors.sigma_d_logit_low, n_un)
+ d_logit_delta[unmeasured] = rng.normal(0.0, sim_priors.sigma_d_logit_delta, n_un)
+ d_log_K[unmeasured] = rng.normal(0.0, sim_priors.sigma_d_log_K, n_un)
+ d_log_n[unmeasured] = rng.normal(0.0, sim_priors.sigma_d_log_n, n_un)
+
+ # Scatter helpers: (1, M) or (1, P) → (1, G)
+ mut_nnz_mut = jnp.array(sim_data.mut_nnz_mut_idx)
+ mut_nnz_geno = jnp.array(sim_data.mut_nnz_geno_idx)
+
+ def _scatter_mut(d_1d):
+ return np.array(
+ apply_mut_matrix(
+ jnp.array(d_1d, dtype=float)[None, :],
+ mut_nnz_mut_idx=mut_nnz_mut,
+ mut_nnz_geno_idx=mut_nnz_geno,
+ num_genotype=G,
+ )
+ )[0]
+
+ # Assemble per-genotype parameters: (G,)
+ logit_low = wt_logit_low + _scatter_mut(d_logit_low)
+ logit_delta = wt_logit_delta + _scatter_mut(d_logit_delta)
+ log_K = wt_log_K + _scatter_mut(d_log_K)
+ log_n = wt_log_n + _scatter_mut(d_log_n)
+
+ # Horseshoe epistasis for pairs (mirrors hill_mut.simulate exactly)
+ if P > 0 and sim_data.pair_nnz_pair_idx is not None:
+ pair_nnz_pair = jnp.array(sim_data.pair_nnz_pair_idx)
+ pair_nnz_geno = jnp.array(sim_data.pair_nnz_geno_idx)
+
+ def _scatter_pair(e_1d):
+ return np.array(
+ apply_pair_matrix(
+ jnp.array(e_1d, dtype=float)[None, :],
+ pair_nnz_pair_idx=pair_nnz_pair,
+ pair_nnz_geno_idx=pair_nnz_geno,
+ num_genotype=G,
+ )
+ )[0]
+
+ tau = float(np.abs(rng.standard_cauchy())) * sim_priors.epi_tau_scale
+
+ if tau > 0.0:
+ c2 = 1.0 / rng.gamma(
+ shape=sim_priors.epi_slab_df / 2.0,
+ scale=2.0 / (sim_priors.epi_slab_df * sim_priors.epi_slab_scale ** 2),
+ )
+
+ def _horseshoe(size):
+ lam = np.abs(rng.standard_cauchy(size))
+ lam_tilde = np.sqrt(c2 * lam ** 2 / (c2 + tau ** 2 * lam ** 2))
+ return rng.standard_normal(size) * tau * lam_tilde
+
+ logit_low += _scatter_pair(_horseshoe(P))
+ logit_delta += _scatter_pair(_horseshoe(P))
+ log_K += _scatter_pair(_horseshoe(P))
+ log_n += _scatter_pair(_horseshoe(P))
+
+ # Convert to theta: (G, C)
+ theta_low_arr = 1.0 / (1.0 + np.exp(-logit_low))
+ theta_high_arr = 1.0 / (1.0 + np.exp(-(logit_low + logit_delta)))
+ hill_n_arr = np.exp(log_n)
+
+ occupancy = 1.0 / (1.0 + np.exp(
+ -hill_n_arr[:, None] * (log_conc[None, :] - log_K[:, None])
+ ))
+ theta_gc = theta_low_arr[:, None] + (theta_high_arr - theta_low_arr)[:, None] * occupancy
+
+ # Build result dict: last occurrence wins for duplicates (matches theta_gc_override semantics)
+ result = {}
+ effective_params = {}
+ for i, g in enumerate(library_genotypes):
+ result[g] = theta_gc[i]
+ effective_params[g] = {
+ "theta_low": float(theta_low_arr[i]),
+ "theta_high": float(theta_high_arr[i]),
+ "log_hill_K": float(log_K[i]),
+ "hill_n": float(hill_n_arr[i]),
+ }
+
+ # Directly-measured multi-mutant entries override the assembled values
+ for g, params in params_dict.items():
+ muts = _parse_muts(g)
+ if len(muts) > 1 and g in result:
+ filled = _fill_params_from_wt(params, wt_params)
+ result[g] = _hill_theta(
+ filled["theta_low"], filled["theta_high"],
+ filled["log_hill_K"], filled["hill_n"],
+ log_conc,
+ )
+ effective_params[g] = {k: filled[k] for k in HILL_PARAM_COLS}
+
+ return result, effective_params
+
+
+# ---------------------------------------------------------------------------
+# Binding theta output
+# ---------------------------------------------------------------------------
+
+def build_binding_theta_from_params(
+ params_dict: dict,
+ binding_concs,
+ titrant_name: str,
+ noise: float,
+ rng: np.random.Generator,
+ wt_params: dict,
+) -> pd.DataFrame:
+ """
+ Compute binding theta rows for measured genotypes.
+
+ Parameters
+ ----------
+ params_dict : dict[str, dict[str, float]]
+ binding_concs : array-like
+ Concentrations (mM) at which binding theta is reported.
+ titrant_name : str
+ noise : float
+ Gaussian noise sigma applied to theta (in [0, 1] space).
+ Set to 0 to disable.
+ rng : np.random.Generator
+ wt_params : dict
+ ``{theta_low, theta_high, log_hill_K, hill_n}`` used to fill NaN params.
+
+ Returns
+ -------
+ pd.DataFrame
+ Columns: ``genotype``, ``titrant_name``, ``titrant_conc``, ``theta_true``.
+ """
+ binding_concs = np.asarray(binding_concs, dtype=float)
+ log_concs = _to_log_conc(binding_concs)
+
+ rows = []
+ for g, params in params_dict.items():
+ filled = _fill_params_from_wt(params, wt_params)
+ theta_vals = _hill_theta(
+ filled["theta_low"], filled["theta_high"],
+ filled["log_hill_K"], filled["hill_n"],
+ log_concs,
+ )
+ for conc, theta in zip(binding_concs, theta_vals):
+ if noise > 0:
+ theta_out = float(np.clip(theta + rng.normal(0.0, noise), 0.0, 1.0))
+ else:
+ theta_out = float(theta)
+ rows.append({
+ "genotype": g,
+ "titrant_name": titrant_name,
+ "titrant_conc": float(conc),
+ "theta_true": theta_out,
+ })
+
+ return pd.DataFrame(rows)
diff --git a/src/tfscreen/simulate/library_prediction.py b/src/tfscreen/simulate/library_prediction.py
index d6651580..df3911d1 100644
--- a/src/tfscreen/simulate/library_prediction.py
+++ b/src/tfscreen/simulate/library_prediction.py
@@ -6,9 +6,21 @@
thermo_to_growth,
)
from tfscreen.simulate.sim_data_class import build_sim_data
+from tfscreen.simulate.sample_theta import sample_theta_stratified, sample_theta_prior
+from tfscreen.simulate.binding_params import (
+ SUPPORTED_COMPONENTS as _BINDING_PARAMS_SUPPORTED,
+ read_binding_genotype_params,
+ build_theta_gc_override_hill_geno,
+ build_theta_gc_override_hill_mut,
+ build_binding_theta_from_params,
+ _wt_params_from_sim_priors,
+)
from tfscreen.genetics import library_manager
+from tfscreen.genetics import standardize_genotypes
import jax
+import numpy as np
+import pandas as pd
from typing import Any, Dict, Union
from pathlib import Path
@@ -39,8 +51,20 @@ def library_prediction(cf: Union[Dict[str, Any], str, Path],
dataframe with predicted fractional occupancy and growth rates for each
of the genotypes in each of the conditions specified in the configuration
genotype_theta_df : pandas.DataFrame
- wide-form dataframe with one row per genotype and one column per unique
- effector concentration giving the ground-truth theta value
+ long-form dataframe with one row per (genotype, titrant_name,
+ titrant_conc) combination; columns: genotype, titrant_name,
+ titrant_conc, theta
+ parameters_df : pandas.DataFrame
+ one row per unique genotype; columns ``dk_geno``, ``activity``, and
+ any scalar per-genotype fields from the theta component (e.g.
+ ``theta_low``, ``theta_high``, ``log_hill_K``, ``hill_n``)
+ binding_theta_df : pandas.DataFrame or None
+ Present when ``binding_data`` is in the config. One row per
+ (genotype, titrant_conc) for the calibration genotypes listed in
+ ``binding_data.genotypes``, with columns ``genotype``,
+ ``titrant_name``, ``titrant_conc``, ``theta_true``. Theta values
+ here are stratified (greedy maximin) across the binding concentrations
+ and are consistent with the growth phenotypes in ``phenotype_df``.
"""
# -------------------------------------------------------------------------
@@ -68,12 +92,172 @@ def library_prediction(cf: Union[Dict[str, Any], str, Path],
thermo_data=cf.get('thermo_data'),
)
- # RNG key for prior-predictive theta sampling
- theta_rng_seed = cf.get('theta_rng_seed', 0)
- theta_rng_key = jax.random.PRNGKey(theta_rng_seed)
+ # Derive both RNG objects from the single seed. They are independent
+ # objects with independent state; using the same integer for both introduces
+ # no correlation between the two streams.
+ seed = cf.get('seed', None)
+ theta_rng_key = jax.random.PRNGKey(seed if seed is not None else 0)
+ rng = np.random.default_rng(seed)
+
+ dk_geno_zero = cf.get('dk_geno_zero', False)
+ if dk_geno_zero:
+ dk_geno_hyper_loc = cf.get('dk_geno_hyper_loc', -3.5)
+ dk_geno_hyper_scale = cf.get('dk_geno_hyper_scale', 1.0)
+ dk_geno_hyper_shift = cf.get('dk_geno_hyper_shift', 0.02)
+ else:
+ dk_geno_hyper_loc = cf['dk_geno_hyper_loc']
+ dk_geno_hyper_scale = cf['dk_geno_hyper_scale']
+ dk_geno_hyper_shift = cf['dk_geno_hyper_shift']
+
+ # -------------------------------------------------------------------------
+ # Stratified theta sampling for binding calibration genotypes
+
+ binding_cfg = cf.get('binding_data')
+ theta_gc_override = {}
+ theta_params_override = {}
+ binding_theta_df = None
+
+ if binding_cfg is not None:
+ binding_concs = binding_cfg['titrant_conc']
+ titrant_name = binding_cfg['titrant_name']
+ binding_noise = binding_cfg.get('noise', 0.0)
+
+ binding_sample_df = pd.DataFrame({"titrant_conc": binding_concs})
+ sorted_binding_concs = np.sort(np.unique(binding_concs))
+ conc_to_col = {float(c): j for j, c in enumerate(sorted_binding_concs)}
+
+ rows = [] # accumulated rows for binding_theta_df
+
+ # --- Section 1: Simulated genotypes (existing 'genotypes' key) ---
+ binding_genotypes = list(standardize_genotypes(binding_cfg.get('genotypes', [])))
+
+ # wt gets its natural unperturbed reference parameters — not a random
+ # draw from the pool. Only non-wt binding genotypes are stratified.
+ non_wt_genotypes = [g for g in binding_genotypes if g != "wt"]
+
+ # wt reference curve (if wt is a simulated binding genotype)
+ wt_binding_gc = None
+ if "wt" in binding_genotypes:
+ single_wt_df = pd.DataFrame({"genotype": ["wt"]})
+ binding_wt_sim = build_sim_data(single_wt_df, binding_sample_df,
+ thermo_data=cf.get('thermo_data'),
+ skip_pairs=True)
+ # Perturbation path for wt (M=0): always returns the fixed sim-priors
+ # reference curve regardless of rng_key, so the binding data matches
+ # the growth simulation exactly.
+ wt_binding_gc, _ = sample_theta_prior(
+ cf['theta_component'], binding_wt_sim, theta_rng_key,
+ sim_priors_overrides=cf.get('theta_sim_priors'),
+ )
+
+ # Stratified curves for non-wt simulated binding genotypes
+ selected_binding_gc = None
+ selected_growth_gc = None
+ if non_wt_genotypes:
+ selected_binding_gc, selected_growth_gc = sample_theta_stratified(
+ component_name=cf['theta_component'],
+ binding_sample_df=binding_sample_df,
+ growth_sample_df=sample_df,
+ rng_key=theta_rng_key,
+ n_select=len(non_wt_genotypes),
+ thermo_data=cf.get('thermo_data'),
+ pool_size=cf.get('binding_stratify_pool_size', 500),
+ priors_overrides=cf.get('theta_priors'),
+ sim_priors_overrides=cf.get('theta_sim_priors'),
+ )
+
+ # Override theta at growth concentrations for non-wt binding genotypes only.
+ # wt is deliberately excluded so thermo_to_growth uses its natural curve.
+ if non_wt_genotypes and selected_growth_gc is not None:
+ theta_gc_override.update({
+ g: selected_growth_gc[i]
+ for i, g in enumerate(non_wt_genotypes)
+ })
+ # Accumulate binding rows for simulated genotypes
+ if "wt" in binding_genotypes and wt_binding_gc is not None:
+ for conc in binding_concs:
+ rows.append({
+ "genotype": "wt",
+ "titrant_name": titrant_name,
+ "titrant_conc": float(conc),
+ "theta_true": float(wt_binding_gc[0, conc_to_col[float(conc)]]),
+ })
+ for i, g in enumerate(non_wt_genotypes):
+ for conc in binding_concs:
+ rows.append({
+ "genotype": g,
+ "titrant_name": titrant_name,
+ "titrant_conc": float(conc),
+ "theta_true": float(selected_binding_gc[i, conc_to_col[float(conc)]]),
+ })
+
+ # --- Section 2: Measured genotype params from CSV ---
+ params_file = binding_cfg.get('genotype_params_file')
+ if params_file is not None:
+ theta_component = cf['theta_component']
+ if theta_component not in _BINDING_PARAMS_SUPPORTED:
+ raise ValueError(
+ f"genotype_params_file is only supported with Hill-based theta "
+ f"components: {sorted(_BINDING_PARAMS_SUPPORTED)}. "
+ f"Got: '{theta_component}'."
+ )
+
+ # Build sim_priors for WT reference / NaN filling
+ from tfscreen.tfmodel.generative.registry import model_registry
+ theta_module = model_registry["theta"][theta_component]
+ sim_params = theta_module.get_sim_hyperparameters()
+ if cf.get('theta_sim_priors'):
+ sim_params.update(cf['theta_sim_priors'])
+ sim_priors = theta_module.SimPriors(**sim_params)
+
+ wt_params = _wt_params_from_sim_priors(sim_priors)
+
+ params_dict = read_binding_genotype_params(params_file)
+
+ # Compute theta overrides for measured genotypes
+ log_conc_growth = np.array(sim_data.log_titrant_conc)
+
+ if theta_component == "hill_geno":
+ measured_override, measured_params = build_theta_gc_override_hill_geno(
+ params_dict, log_conc_growth, wt_params
+ )
+ else: # hill_mut
+ measured_override, measured_params = build_theta_gc_override_hill_mut(
+ params_dict=params_dict,
+ library_genotypes=library_df["genotype"].tolist(),
+ sim_data=sim_data,
+ sim_priors=sim_priors,
+ log_conc=log_conc_growth,
+ rng=rng,
+ )
+
+ # Measured data takes precedence over stratified simulated data
+ theta_gc_override.update(measured_override)
+ theta_params_override.update(measured_params)
+
+ # Accumulate binding rows for measured genotypes (overrides simulated rows)
+ # Store noise-free theta_true; noise is applied later by simulate_cli
+ # _generate_binding_data, consistent with the simulated-genotypes path.
+ measured_binding_df = build_binding_theta_from_params(
+ params_dict=params_dict,
+ binding_concs=binding_concs,
+ titrant_name=titrant_name,
+ noise=0.0,
+ rng=rng,
+ wt_params=wt_params,
+ )
+ # Remove any existing rows for genotypes now covered by measured data
+ measured_genotypes = set(params_dict.keys())
+ rows = [r for r in rows if r["genotype"] not in measured_genotypes]
+ rows.extend(measured_binding_df.to_dict("records"))
+
+ binding_theta_df = pd.DataFrame(rows) if rows else None
+
+ # -------------------------------------------------------------------------
# Calculate phenotype for each genotype across all conditions in sample_df
- phenotype_df, genotype_theta_df = thermo_to_growth(
+
+ phenotype_df, genotype_theta_df, parameters_df = thermo_to_growth(
genotypes=library_df["genotype"],
sim_data=sim_data,
sample_df=sample_df,
@@ -81,14 +265,19 @@ def library_prediction(cf: Union[Dict[str, Any], str, Path],
theta_rng_key=theta_rng_key,
growth_params=cf['growth'],
theta_priors_overrides=cf.get('theta_priors'),
- dk_geno_hyper_loc=cf['dk_geno_hyper_loc'],
- dk_geno_hyper_scale=cf['dk_geno_hyper_scale'],
- dk_geno_hyper_shift=cf['dk_geno_hyper_shift'],
+ theta_sim_priors_overrides=cf.get('theta_sim_priors'),
+ dk_geno_hyper_loc=dk_geno_hyper_loc,
+ dk_geno_hyper_scale=dk_geno_hyper_scale,
+ dk_geno_hyper_shift=dk_geno_hyper_shift,
+ dk_geno_zero=dk_geno_zero,
activity_wt=cf.get('activity_wt', 1.0),
activity_mut_scale=cf.get('activity_mut_scale', 0.0),
+ rng=rng,
activity_component=cf.get('activity_component', 'fixed'),
activity_priors_overrides=cf.get('activity_priors'),
theta_rescale=cf.get('theta_rescale', 'passthrough'),
+ theta_gc_override=theta_gc_override,
+ theta_params_override=theta_params_override or None,
)
- return library_df, phenotype_df, genotype_theta_df
+ return library_df, phenotype_df, genotype_theta_df, parameters_df, binding_theta_df
diff --git a/src/tfscreen/simulate/run_simulation.py b/src/tfscreen/simulate/run_simulation.py
index 90e0b24b..374210f1 100644
--- a/src/tfscreen/simulate/run_simulation.py
+++ b/src/tfscreen/simulate/run_simulation.py
@@ -19,7 +19,7 @@ def _setup_file_output(output_dir,
err = "if output_dir is specified, output_prefix must be a string\n"
raise ValueError(err)
- roots = ["library","phenotype","genotype_theta","sample","counts"]
+ roots = ["library","parameters","genotype_theta","sample","counts"]
files = [f"{output_prefix}{r}.csv" for r in roots]
files_with_path = [os.path.join(output_dir,f) for f in files]
file_dict = dict([(roots[i],files_with_path[i]) for i in range(len(roots))])
@@ -82,14 +82,14 @@ def run_simulation(cf: Union[Dict[str, Any], str, Path],
file_dict = _setup_file_output(output_dir,output_prefix)
# Build library and predict its phenotypes
- library_df, phenotype_df, genotype_theta_df = library_prediction(cf)
+ library_df, phenotype_df, genotype_theta_df, parameters_df, _ = library_prediction(cf)
# Perform selection experiment
sample_df, counts_df = selection_experiment(cf,library_df,phenotype_df)
# Prepare outputs
out_dict = {"library":library_df,
- "phenotype":phenotype_df,
+ "parameters":parameters_df,
"genotype_theta":genotype_theta_df,
"sample":sample_df,
"counts":counts_df}
diff --git a/src/tfscreen/simulate/sample_theta.py b/src/tfscreen/simulate/sample_theta.py
index 4f112716..43398531 100644
--- a/src/tfscreen/simulate/sample_theta.py
+++ b/src/tfscreen/simulate/sample_theta.py
@@ -6,12 +6,24 @@
path: instead of reading fixed thermodynamic parameters from a file, we draw
parameters from the model prior, making the simulation consistent with the
model family being inferred.
+
+sample_theta_stratified draws a large pool of independent theta curves from
+the prior at the binding concentrations, selects n_select curves that
+maximally span the theta-curve space (greedy maximin), then returns the
+selected curves at both binding and growth concentrations. This ensures
+calibration genotypes cover diverse binding behaviors rather than clustering
+by chance.
"""
+import inspect
import numpy as np
+import pandas as pd
+import tqdm
+import jax
from numpyro import handlers
from tfscreen.tfmodel.generative.registry import model_registry
+from tfscreen.simulate.sim_data_class import build_sim_data
# Components excluded from simulation (calibration-only, no generative prior).
_EXCLUDED = frozenset({"_simple"})
@@ -20,36 +32,53 @@
def sample_theta_prior(component_name,
sim_data,
rng_key,
- priors_overrides=None):
+ priors_overrides=None,
+ sim_priors_overrides=None,
+ force_prior_predictive=False):
"""
- Draw one sample of theta from the prior of a registered theta component.
+ Draw one sample of theta from a registered theta component.
+
+ If the component defines a ``simulate`` function (detected via
+ ``inspect.isfunction``) and ``force_prior_predictive`` is False, the
+ perturbation-based path is used: wildtype reference parameters are drawn
+ from ``SimPriors`` built from ``get_sim_hyperparameters()`` plus any
+ ``sim_priors_overrides``.
+
+ Otherwise the prior-predictive path is used: ``define_model`` is called
+ inside a seeded NumPyro trace, and ``run_model`` produces ``theta_gc``.
+ Set ``force_prior_predictive=True`` to always use this path — this is
+ required for mutation-decomposed components (e.g. ``hill_mut``) when the
+ pool library has only the wt genotype and M=0 mutations, because the
+ perturbation path would produce zero deltas and a homogeneous pool.
Parameters
----------
component_name : str
- A key in ``model_registry["theta"]`` (e.g. ``"hill"``,
- ``"mwc_dimer_lnK_mut"``). ``"simple"`` is not supported.
+ A key in ``model_registry["theta"]`` (e.g. ``"hill_geno"``,
+ ``"mwc_dimer_lnK_mut"``). ``"_simple"`` is not supported.
sim_data : SimData
Built by ``build_sim_data``. Must have all fields required by the
chosen component.
rng_key : jax.random.PRNGKey
- Seed for the NumPyro sampler.
+ Seed for sampling (NumPyro or NumPy depending on the path taken).
priors_overrides : dict or None
- Key-value overrides applied to ``get_hyperparameters()`` before
- constructing the priors object. Use to set physical constants such
- as ``theta_tf_total_M`` or concentration unit scales without
- modifying the component defaults.
+ Overrides for the inference ``ModelPriors`` (prior-predictive path
+ only). Ignored when the component provides ``simulate``.
+ sim_priors_overrides : dict or None
+ Overrides for ``SimPriors`` (perturbation path only). Ignored when
+ the component does not provide ``simulate``.
+ force_prior_predictive : bool, default False
+ When True, always use the prior-predictive path (``define_model`` +
+ ``handlers.seed``), even if the component provides ``simulate``.
Returns
-------
theta : np.ndarray, shape (num_genotype, num_titrant_conc)
Ground-truth fractional occupancy for each genotype at each unique
- effector concentration. Store alongside phenotype_df to enable
- comparison with fitted posteriors.
+ effector concentration.
theta_param : ThetaParam
- Sampled parameter pytree (structure is model-specific). Contains
- the underlying equilibrium constants, Hill parameters, etc. that
- generated ``theta``.
+ Sampled parameter pytree. When the perturbation path is used,
+ ``mu`` and ``sigma`` fields are ``None``.
Raises
------
@@ -71,17 +100,24 @@ def sample_theta_prior(component_name,
module = theta_registry[component_name]
- # Build priors: start from module defaults, then apply any overrides.
- params = module.get_hyperparameters()
- if priors_overrides:
- params.update(priors_overrides)
- priors = module.ModelPriors(**params)
+ # Perturbation-based path: component defines a simulate() function.
+ if not force_prior_predictive and inspect.isfunction(getattr(module, "simulate", None)):
+ sim_params = module.get_sim_hyperparameters()
+ if sim_priors_overrides:
+ sim_params.update(sim_priors_overrides)
+ sim_priors = module.SimPriors(**sim_params)
+ return module.simulate("theta", sim_data, sim_priors, rng_key)
- # Seed the NumPyro trace and call define_model directly.
+ # Prior-predictive fallback: seed the NumPyro trace and call define_model.
# handlers.seed gives every pyro.sample site a reproducible draw
# without needing a full Predictive wrapper. pyro.param sites return
# their init_value (module default) when no param_store is active,
# which is the correct behaviour for prior-predictive simulation.
+ params = module.get_hyperparameters()
+ if priors_overrides:
+ params.update(priors_overrides)
+ priors = module.ModelPriors(**params)
+
with handlers.seed(rng_seed=rng_key):
theta_param = module.define_model("theta", sim_data, priors)
@@ -90,3 +126,149 @@ def sample_theta_prior(component_name,
theta_gc = np.array(theta_tcg[0]).T # (G, C)
return theta_gc, theta_param
+
+
+def _greedy_maximin(theta_gc, n_select):
+ """
+ Select n_select rows from theta_gc that maximally span the theta-curve space.
+
+ Uses greedy farthest-point selection (maximin): starts with the row of
+ greatest dynamic range (max - min across concentrations), then repeatedly
+ adds the row that maximises the minimum Euclidean distance to the
+ already-selected set.
+
+ Parameters
+ ----------
+ theta_gc : np.ndarray, shape (pool_size, n_conc)
+ n_select : int
+
+ Returns
+ -------
+ np.ndarray of int, shape (n_select,)
+ Indices into theta_gc of the selected rows, in selection order.
+ """
+ pool_size = theta_gc.shape[0]
+ if n_select >= pool_size:
+ return np.arange(pool_size, dtype=int)
+
+ theta_range = theta_gc.max(axis=1) - theta_gc.min(axis=1)
+ selected = [int(np.argmax(theta_range))]
+
+ while len(selected) < n_select:
+ selected_arr = theta_gc[selected] # (k, C)
+ # Squared distance from every pool member to its nearest selected point
+ diff = theta_gc[:, None, :] - selected_arr[None, :, :] # (pool, k, C)
+ sq_dist = (diff ** 2).sum(axis=-1) # (pool, k)
+ min_sq_dist = sq_dist.min(axis=1) # (pool,)
+ min_sq_dist[selected] = -1.0 # exclude already selected
+ selected.append(int(np.argmax(min_sq_dist)))
+
+ return np.array(selected, dtype=int)
+
+
+def sample_theta_stratified(component_name,
+ binding_sample_df,
+ growth_sample_df,
+ rng_key,
+ n_select,
+ thermo_data=None,
+ pool_size=500,
+ priors_overrides=None,
+ sim_priors_overrides=None):
+ """
+ Sample a pool of theta curves from the prior and select n_select that
+ maximally span the theta-curve space (greedy maximin on binding concentrations).
+
+ Builds an internal pool library of pool_size "wt" genotypes. For per-genotype
+ theta components (e.g. hill_geno), each pool member receives an independent
+ i.i.d. draw from the prior. Mutation-decomposed components are not supported.
+
+ The same JAX rng_key is used for both the binding-concentration and
+ growth-concentration sampling calls. Because per-genotype parameter draws
+ in the NumPyro trace depend only on plate sizes (not on concentration values),
+ both calls produce the same underlying parameter sets evaluated at their
+ respective concentrations, so the returned curves are mutually consistent.
+
+ Parameters
+ ----------
+ component_name : str
+ A key in model_registry["theta"].
+ binding_sample_df : pd.DataFrame
+ Must contain a "titrant_conc" column (mM). These concentrations are
+ used to select which parameter sets maximally span the binding space.
+ growth_sample_df : pd.DataFrame
+ Must contain a "titrant_conc" column (mM). The selected parameter sets
+ are also evaluated here so they can be injected into the growth simulation.
+ rng_key : jax.random.PRNGKey
+ n_select : int
+ Number of genotypes to select (= number of calibration binding genotypes).
+ thermo_data : str or None
+ Forwarded to build_sim_data.
+ pool_size : int, default 500
+ Number of candidate parameter sets to draw before selection.
+ priors_overrides : dict or None
+ Forwarded to sample_theta_prior (prior-predictive path).
+ sim_priors_overrides : dict or None
+ Forwarded to sample_theta_prior (perturbation path).
+
+ Returns
+ -------
+ binding_theta_gc : np.ndarray, shape (n_select, n_binding_concs)
+ Selected theta curves at binding concentrations (sorted ascending).
+ growth_theta_gc : np.ndarray, shape (n_select, n_growth_concs)
+ Same selected parameter sets evaluated at growth concentrations.
+
+ Raises
+ ------
+ ValueError
+ If n_select > pool_size.
+ """
+ if n_select > pool_size:
+ raise ValueError(
+ f"n_select ({n_select}) must not exceed pool_size ({pool_size}). "
+ f"Increase pool_size or reduce the number of binding genotypes."
+ )
+
+ if isinstance(rng_key, int):
+ rng_key = jax.random.PRNGKey(rng_key)
+
+ # Build sim_data for a single genotype at each concentration set.
+ # The pool is built by calling sample_theta_prior once per member with a
+ # unique key derived via fold_in. This gives each member its own
+ # hyperprior draw (loc, scale, …) rather than sharing one draw across all
+ # pool members — which would concentrate the pool around a single binding
+ # regime regardless of pool_size.
+ single_library_df = pd.DataFrame({"genotype": ["wt"]})
+ binding_sim_data = build_sim_data(single_library_df, binding_sample_df,
+ thermo_data=thermo_data, skip_pairs=True)
+ growth_sim_data = build_sim_data(single_library_df, growth_sample_df,
+ thermo_data=thermo_data, skip_pairs=True)
+
+ # Always use the prior-predictive path for pool building. The perturbation
+ # path (simulate()) requires mutations to generate diversity; with a wt-only
+ # single-genotype pool library (M=0), it produces zero deltas and an
+ # entirely homogeneous pool regardless of pool_size. The prior-predictive
+ # path samples WT-level parameters (log_K_wt, logit_low_wt, …) from their
+ # broad Normal priors, giving real diversity across pool members.
+ print(f"Building stratified pool ({pool_size} candidates)... ",
+ end="", flush=True)
+ pool_binding_rows = []
+ pool_growth_rows = []
+ for i in tqdm.tqdm(range(pool_size)):
+ key_i = jax.random.fold_in(rng_key, i)
+ theta_b, _ = sample_theta_prior(component_name, binding_sim_data, key_i,
+ priors_overrides, sim_priors_overrides,
+ force_prior_predictive=True)
+ theta_g, _ = sample_theta_prior(component_name, growth_sim_data, key_i,
+ priors_overrides, sim_priors_overrides,
+ force_prior_predictive=True)
+ pool_binding_rows.append(theta_b[0]) # (C_b,) — single genotype
+ pool_growth_rows.append(theta_g[0]) # (C_g,)
+ print("Done.", flush=True)
+
+ pool_binding_gc = np.stack(pool_binding_rows) # (pool_size, C_b)
+ pool_growth_gc = np.stack(pool_growth_rows) # (pool_size, C_g)
+
+ selected_indices = _greedy_maximin(pool_binding_gc, n_select)
+
+ return pool_binding_gc[selected_indices], pool_growth_gc[selected_indices]
diff --git a/src/tfscreen/simulate/scripts/run_simulation_cli.py b/src/tfscreen/simulate/scripts/run_simulation_cli.py
deleted file mode 100644
index 4fff97a7..00000000
--- a/src/tfscreen/simulate/scripts/run_simulation_cli.py
+++ /dev/null
@@ -1,208 +0,0 @@
-import os
-import numpy as np
-import pandas as pd
-import jax
-
-import tfscreen
-from tfscreen.simulate import library_prediction, selection_experiment
-from tfscreen.simulate.sim_data_class import build_sim_data
-from tfscreen.simulate.sample_theta import sample_theta_prior
-from tfscreen.genetics import standardize_genotypes
-from tfscreen.process_raw import counts_to_lncfu
-from tfscreen.util.cli.generalized_main import generalized_main
-
-
-def _generate_binding_data(cf, library_df, binding_cfg, rng):
- """
- Generate simulated binding curve data for specific genotypes.
-
- Samples theta from the same prior used by library_prediction (same
- theta_component and theta_rng_seed) but evaluated at the binding
- concentrations. Adds Gaussian noise to the ground-truth theta values.
-
- Parameters
- ----------
- cf : dict
- Full run configuration (already read from YAML).
- library_df : pandas.DataFrame
- Library returned by library_prediction (has a 'genotype' column).
- binding_cfg : dict
- The 'binding_data' sub-dict from the config. Must contain:
- genotypes : list of genotype strings
- titrant_name: str, name of the titrant (e.g. 'iptg')
- titrant_conc: list of concentrations (mM)
- noise : float, sigma for Gaussian noise on theta_obs
- rng : numpy.random.Generator
-
- Returns
- -------
- pandas.DataFrame
- Columns: genotype, titrant_name, titrant_conc, theta_obs, theta_std
- """
- genotypes = list(standardize_genotypes(binding_cfg["genotypes"]))
- titrant_name = binding_cfg["titrant_name"]
- titrant_conc = list(binding_cfg["titrant_conc"])
- noise = float(binding_cfg.get("noise", 0.0))
-
- # Build a SimData for just the binding concentrations
- binding_sample_df = pd.DataFrame({
- "titrant_conc": titrant_conc,
- "condition_pre": "binding",
- "condition_sel": "binding",
- })
- sim_data = build_sim_data(
- library_df=library_df,
- sample_df=binding_sample_df,
- thermo_data=cf.get("thermo_data"),
- )
-
- # Sample theta from the prior using the same seed as the main run
- theta_rng_key = jax.random.PRNGKey(cf.get("theta_rng_seed", 0))
- theta_gc, _ = sample_theta_prior(
- component_name=cf["theta_component"],
- sim_data=sim_data,
- rng_key=theta_rng_key,
- priors_overrides=cf.get("theta_priors"),
- )
-
- # Build lookup: genotype string → row index in library_df (sim_data order)
- all_genotypes = library_df["genotype"].tolist()
- geno_to_idx = {g: i for i, g in enumerate(all_genotypes)}
-
- # Sorted unique concentrations (same order as sim_data.titrant_conc)
- sorted_concs = np.sort(np.unique(titrant_conc))
- conc_to_col = {float(c): i for i, c in enumerate(sorted_concs)}
-
- rows = []
- for g in genotypes:
- if g not in geno_to_idx:
- raise ValueError(
- f"Genotype '{g}' in binding_data.genotypes is not in the library."
- )
- row_idx = geno_to_idx[g]
- for conc in titrant_conc:
- col_idx = conc_to_col[float(conc)]
- theta_true = float(theta_gc[row_idx, col_idx])
- if noise > 0:
- theta_obs = float(np.clip(theta_true + rng.normal(0, noise), 0, 1))
- else:
- theta_obs = theta_true
- rows.append({
- "genotype": g,
- "titrant_name": titrant_name,
- "titrant_conc": conc,
- "theta_obs": theta_obs,
- "theta_std": noise,
- })
-
- return pd.DataFrame(rows)
-
-
-def run_simulation_from_config(
- config_file,
- output_dir,
- output_prefix="tfscreen_",
- num_replicates=2,
-):
- """
- Simulate a TF selection experiment from a YAML configuration file.
-
- Runs library_prediction once to establish ground-truth phenotypes, then
- simulates num_replicates independent experimental replicates using
- selection_experiment. Writes library, phenotype, genotype_theta, and
- analysis-ready growth CSV files. If the config contains a 'binding_data'
- block, also writes a simulated binding curve CSV.
-
- Parameters
- ----------
- config_file : str
- Path to the YAML run configuration file.
- output_dir : str
- Directory to write output CSV files into (created if absent).
- output_prefix : str
- Prefix for all output filenames. Default 'tfscreen_'.
- num_replicates : int
- Number of independent experimental replicates to simulate. Default 2.
- """
- cf = tfscreen.util.read_yaml(config_file)
-
- os.makedirs(output_dir, exist_ok=True)
-
- def out_path(name):
- return os.path.join(output_dir, f"{output_prefix}{name}.csv")
-
- output_names = ["library", "phenotype", "genotype_theta", "growth"]
- if "binding_data" in cf:
- output_names.append("binding")
-
- existing = [n for n in output_names if os.path.exists(out_path(n))]
- if existing:
- paths = ", ".join(out_path(n) for n in existing)
- raise FileExistsError(f"Output files already exist: {paths}")
-
- base_seed = cf.get("random_seed", None)
- rng = np.random.default_rng(base_seed)
-
- # -------------------------------------------------------------------------
- # Ground-truth library and phenotypes (deterministic given the config)
-
- library_df, phenotype_df, genotype_theta_df = library_prediction(cf)
-
- # -------------------------------------------------------------------------
- # Simulate independent replicates
-
- all_sample_parts = []
- all_counts_parts = []
- sample_id_offset = 0
-
- for rep in range(1, num_replicates + 1):
- print(f"\n--- Replicate {rep} of {num_replicates} ---", flush=True)
-
- # Give each replicate a distinct (but reproducible) random seed so
- # that replicates differ even when a base seed is set.
- rep_cf = dict(cf)
- rep_cf["random_seed"] = (
- base_seed * num_replicates + rep if base_seed is not None else None
- )
-
- rep_phenotype_df = phenotype_df.copy()
- rep_phenotype_df["replicate"] = rep
-
- sample_df_rep, counts_df_rep = selection_experiment(
- rep_cf, library_df, rep_phenotype_df
- )
-
- # Shift sample IDs so they are globally unique across replicates
- max_id = int(sample_df_rep.index.max()) + 1
- sample_df_rep = sample_df_rep.copy()
- counts_df_rep = counts_df_rep.copy()
- sample_df_rep.index = sample_df_rep.index + sample_id_offset
- sample_df_rep["sample"] = sample_df_rep.index
- counts_df_rep["sample"] = counts_df_rep["sample"] + sample_id_offset
-
- all_sample_parts.append(sample_df_rep)
- all_counts_parts.append(counts_df_rep)
- sample_id_offset += max_id
-
- combined_sample_df = pd.concat(all_sample_parts)
- combined_counts_df = pd.concat(all_counts_parts, ignore_index=True)
-
- growth_df = counts_to_lncfu(combined_sample_df, combined_counts_df)
-
- # -------------------------------------------------------------------------
- # Write outputs
-
- library_df.to_csv(out_path("library"), index=False)
- phenotype_df.to_csv(out_path("phenotype"), index=False)
- genotype_theta_df.to_csv(out_path("genotype_theta"), index=False)
- growth_df.to_csv(out_path("growth"), index=False)
- print(f"\nWrote: {', '.join(out_path(n) for n in ['library', 'phenotype', 'genotype_theta', 'growth'])}")
-
- if "binding_data" in cf:
- binding_df = _generate_binding_data(cf, library_df, cf["binding_data"], rng)
- binding_df.to_csv(out_path("binding"), index=False)
- print(f"Wrote: {out_path('binding')}")
-
-
-def main():
- return generalized_main(run_simulation_from_config)
diff --git a/src/tfscreen/simulate/scripts/simulate_cli.py b/src/tfscreen/simulate/scripts/simulate_cli.py
new file mode 100644
index 00000000..713c2809
--- /dev/null
+++ b/src/tfscreen/simulate/scripts/simulate_cli.py
@@ -0,0 +1,345 @@
+import os
+import yaml
+import numpy as np
+import pandas as pd
+
+import tfscreen
+from tfscreen.simulate import library_prediction, selection_experiment
+from tfscreen.simulate.selection_experiment import _sim_index_hop
+from tfscreen.genetics import standardize_genotypes
+from tfscreen.process_raw import counts_to_lncfu
+from tfscreen.util.cli.generalized_main import generalized_main
+
+
+def _generate_binding_data(binding_cfg, rng, binding_theta_df):
+ """
+ Generate simulated binding curve data for specific genotypes.
+
+ Uses pre-computed (stratified) theta values from library_prediction and
+ adds Gaussian noise to produce observed theta values.
+
+ Parameters
+ ----------
+ binding_cfg : dict
+ The 'binding_data' sub-dict from the config. Must contain:
+ genotypes : list of genotype strings
+ titrant_name: str, name of the titrant (e.g. 'iptg')
+ titrant_conc: list of concentrations (mM)
+ noise : float, sigma for Gaussian noise on theta_obs
+ rng : numpy.random.Generator
+ binding_theta_df : pandas.DataFrame
+ Pre-computed binding theta from library_prediction. Must contain
+ columns ``genotype``, ``titrant_conc``, ``theta_true``.
+
+ Returns
+ -------
+ pandas.DataFrame
+ Columns: genotype, titrant_name, titrant_conc, theta_obs, theta_std
+ """
+ titrant_name = binding_cfg["titrant_name"]
+ titrant_conc = list(binding_cfg["titrant_conc"])
+ noise = float(binding_cfg.get("noise", 0.0))
+
+ # Build lookup: (genotype, conc) → theta_true
+ theta_lookup = {
+ (row["genotype"], float(row["titrant_conc"])): row["theta_true"]
+ for _, row in binding_theta_df.iterrows()
+ }
+
+ # Genotypes to output: explicit list from config, or every genotype in binding_theta_df.
+ raw_genotypes = binding_cfg.get("genotypes")
+ if raw_genotypes:
+ genotypes = list(standardize_genotypes(raw_genotypes))
+ else:
+ genotypes = list(binding_theta_df["genotype"].unique())
+
+ rows = []
+ for g in genotypes:
+ for conc in titrant_conc:
+ key = (g, float(conc))
+ if key not in theta_lookup:
+ raise ValueError(
+ f"No pre-computed theta for genotype '{g}' at conc {conc}. "
+ f"Ensure binding_data genotypes and titrant_conc match the config."
+ )
+ theta_true = float(theta_lookup[key])
+ if noise > 0:
+ theta_obs = float(np.clip(theta_true + rng.normal(0, noise), 0, 1))
+ else:
+ theta_obs = theta_true
+ rows.append({
+ "genotype": g,
+ "titrant_name": titrant_name,
+ "titrant_conc": conc,
+ "theta_obs": theta_obs,
+ "theta_std": noise,
+ })
+
+ return pd.DataFrame(rows)
+
+
+def _generate_presplit_data(combined_sample_df, combined_counts_df, cf, rng):
+ """
+ Generate simulated pre-split (t = -t_pre) data from selection-experiment
+ outputs.
+
+ For each unique ``(replicate, condition_pre)`` pair, this function draws a
+ synthetic sequencing sample from the initial library-frequency distribution
+ (encoded as ``ln_cfu_0`` in ``combined_counts_df``). The resulting counts
+ are converted to ``ln_cfu`` and ``ln_cfu_std`` using the same
+ variance-propagation formula as ``counts_to_lncfu``, yielding a DataFrame
+ that can be passed directly to ``tfs-configure-model --presplit_df``.
+
+ Parameters
+ ----------
+ combined_sample_df : pandas.DataFrame
+ Concatenated sample metadata across all replicates (index = sample ID).
+ Must contain columns ``replicate`` and ``condition_pre``.
+ combined_counts_df : pandas.DataFrame
+ Concatenated genotype counts across all replicates. Must contain
+ columns ``sample``, ``genotype``, and ``ln_cfu_0`` (the ground-truth
+ initial log-CFU per genotype computed by the simulation).
+ cf : dict
+ Full run configuration (already read from YAML). Used keys:
+ ``cfu0``, ``total_num_reads``, ``prob_index_hop``. The optional ``presplit_data`` sub-dict may
+ contain a scalar ``noise`` key (default 0) to add extra Gaussian noise
+ to ``ln_cfu`` on the log scale.
+ rng : numpy.random.Generator
+ Seeded random-number generator shared with the rest of the simulation.
+
+ Returns
+ -------
+ pandas.DataFrame
+ Columns: ``replicate``, ``condition_pre``, ``genotype``,
+ ``ln_cfu``, ``ln_cfu_std``, ``ln_cfu_0_true``.
+ ``ln_cfu_0_true`` records the simulation ground truth for validation.
+ Rows with zero initial frequency (genotypes absent from the
+ transformation pool) receive ``ln_cfu = NaN``.
+ """
+ total_cfu0 = float(cf["cfu0"])
+ presplit_cfg = cf.get("presplit_data", {})
+ extra_noise = float(presplit_cfg.get("noise", 0.0)) if presplit_cfg else 0.0
+ pseudocount = 1
+
+ # Reads per presplit sample: use the same budget as the main experiment
+ # (total reads / total selection samples).
+ total_num_reads = int(cf["total_num_reads"])
+ total_num_samples = len(combined_sample_df)
+ reads_per_sample = max(1, int(round(total_num_reads / total_num_samples)))
+
+ # Attach condition_pre and replicate to each (genotype, sample) row so we
+ # can group by (replicate, condition_pre) below.
+ sample_meta = (combined_sample_df
+ .reset_index()[["sample", "replicate", "condition_pre"]]
+ .drop_duplicates("sample"))
+ counts_meta = pd.merge(combined_counts_df, sample_meta, on="sample", how="left")
+
+ # One row per (replicate, condition_pre, genotype), keeping the first
+ # occurrence of ln_cfu_0 (same value for all selection samples within a
+ # library group).
+ source = (counts_meta
+ .groupby(["replicate", "condition_pre", "genotype"], observed=True)
+ .first()
+ .reset_index()[["replicate", "condition_pre", "genotype", "ln_cfu_0"]])
+
+ rows = []
+ for (rep, cp), grp in source.groupby(["replicate", "condition_pre"],
+ observed=True):
+ genos = grp["genotype"].values
+ ln_cfu0 = grp["ln_cfu_0"].values
+
+ # Convert ground-truth ln_cfu_0 to initial frequencies.
+ # Genotypes with ln_cfu_0 = -inf (absent from transformation pool)
+ # get frequency 0.
+ cfu0_raw = np.where(np.isfinite(ln_cfu0), np.exp(ln_cfu0), 0.0)
+ total_cfu_group = cfu0_raw.sum()
+ if total_cfu_group == 0:
+ continue
+ freqs = cfu0_raw / total_cfu_group
+
+ # Simulate multinomial sequencing draw from the initial distribution.
+ counts = rng.multinomial(reads_per_sample, freqs)
+
+ # Optionally apply index hopping (same as the selection samples).
+ counts = _sim_index_hop(counts, cf.get("prob_index_hop"), rng)
+
+ # ---------- counts → ln_cfu (mirrors counts_to_lncfu logic) ----------
+ sample_cfu = total_cfu0
+
+ total_adjusted = counts.sum() + len(counts) * pseudocount
+ adj_counts = counts + pseudocount
+ freq_est = adj_counts / total_adjusted
+ cfu_est = freq_est * sample_cfu
+
+ # Variance from binomial frequency uncertainty only
+ var_freq = freq_est * (1.0 - freq_est) / total_adjusted
+ with np.errstate(divide="ignore", invalid="ignore"):
+ rel_var_freq = np.where(freq_est > 0,
+ var_freq / (freq_est ** 2), 0.0)
+ ln_cfu_var = rel_var_freq
+
+ # Optional additional noise on the log scale
+ if extra_noise > 0.0:
+ ln_cfu_var = ln_cfu_var + extra_noise ** 2
+ ln_cfu_shift = rng.normal(0.0, extra_noise, size=len(genos))
+ else:
+ ln_cfu_shift = np.zeros(len(genos))
+
+ with np.errstate(divide="ignore"):
+ ln_cfu_vals = np.log(cfu_est) + ln_cfu_shift
+
+ # Zero-count entries become NaN
+ ln_cfu_vals = np.where(cfu_est > 0, ln_cfu_vals, np.nan)
+ ln_cfu_var = np.where(cfu_est > 0, ln_cfu_var, np.nan)
+ ln_cfu_std_vals = np.sqrt(ln_cfu_var)
+
+ for i, geno in enumerate(genos):
+ rows.append({
+ "replicate": rep,
+ "condition_pre": cp,
+ "genotype": geno,
+ "ln_cfu": float(ln_cfu_vals[i]),
+ "ln_cfu_std": float(ln_cfu_std_vals[i]),
+ "ln_cfu_0_true": float(ln_cfu0[i]),
+ })
+
+ return pd.DataFrame(rows)
+
+
+def run_simulation_from_config(
+ config_file,
+ output_dir,
+ output_prefix="tfs_sim_",
+ num_replicates=2,
+ seed=None,
+):
+ """
+ Simulate a TF selection experiment from a YAML configuration file.
+
+ Runs library_prediction once to establish ground-truth phenotypes, then
+ simulates num_replicates independent experimental replicates using
+ selection_experiment. Writes library, parameters, genotype_theta (long-form:
+ genotype/titrant_name/titrant_conc/theta), and analysis-ready growth CSV
+ files. If the config contains a 'binding_data' block, also writes a
+ simulated binding curve CSV.
+
+ Parameters
+ ----------
+ config_file : str
+ Path to the YAML run configuration file.
+ output_dir : str
+ Directory to write output CSV files into (created if absent).
+ output_prefix : str
+ Prefix for all output filenames. Default 'tfs_sim_'.
+ num_replicates : int
+ Number of independent experimental replicates to simulate. Default 2.
+ seed : int, optional
+ Random seed. Overrides seed in the config file when provided.
+ """
+ cf = tfscreen.util.read_yaml(config_file)
+ if seed is not None:
+ cf["seed"] = seed
+
+ os.makedirs(output_dir, exist_ok=True)
+
+ def out_path(name):
+ return os.path.join(output_dir, f"{output_prefix}{name}.csv")
+
+ config_out = os.path.join(output_dir, f"{output_prefix}input-config.yaml")
+
+ output_names = ["library", "parameters", "genotype_theta", "growth"]
+ if "binding_data" in cf:
+ output_names.append("binding")
+ if "presplit_data" in cf:
+ output_names.append("presplit")
+
+ existing = [out_path(n) for n in output_names if os.path.exists(out_path(n))]
+ if os.path.exists(config_out):
+ existing.append(config_out)
+ if existing:
+ paths = ", ".join(existing)
+ raise FileExistsError(
+ f"Output files already exist: {paths}\n"
+ f"Delete them or choose a different output_dir / output_prefix "
+ f"before re-running."
+ )
+
+ base_seed = cf.get("seed", None)
+ rng = np.random.default_rng(base_seed)
+
+ # -------------------------------------------------------------------------
+ # Ground-truth library and phenotypes (deterministic given the config)
+
+ library_df, phenotype_df, genotype_theta_df, parameters_df, binding_theta_df = library_prediction(cf)
+
+ # -------------------------------------------------------------------------
+ # Simulate independent replicates
+
+ all_sample_parts = []
+ all_counts_parts = []
+ sample_id_offset = 0
+
+ for rep in range(1, num_replicates + 1):
+ print(f"\n--- Replicate {rep} of {num_replicates} ---", flush=True)
+
+ # Give each replicate a distinct (but reproducible) random seed so
+ # that replicates differ even when a base seed is set.
+ rep_cf = dict(cf)
+ rep_cf["seed"] = (
+ base_seed * num_replicates + rep if base_seed is not None else None
+ )
+
+ rep_phenotype_df = phenotype_df.copy()
+ rep_phenotype_df["replicate"] = rep
+
+ sample_df_rep, counts_df_rep = selection_experiment(
+ rep_cf, library_df, rep_phenotype_df
+ )
+
+ # Shift sample IDs so they are globally unique across replicates
+ max_id = int(sample_df_rep.index.max()) + 1
+ sample_df_rep = sample_df_rep.copy()
+ counts_df_rep = counts_df_rep.copy()
+ sample_df_rep.index = sample_df_rep.index + sample_id_offset
+ sample_df_rep["sample"] = sample_df_rep.index
+ counts_df_rep["sample"] = counts_df_rep["sample"] + sample_id_offset
+
+ all_sample_parts.append(sample_df_rep)
+ all_counts_parts.append(counts_df_rep)
+ sample_id_offset += max_id
+
+ combined_sample_df = pd.concat(all_sample_parts)
+ combined_counts_df = pd.concat(all_counts_parts, ignore_index=True)
+
+ growth_df = counts_to_lncfu(combined_sample_df, combined_counts_df)
+
+ # -------------------------------------------------------------------------
+ # Write outputs
+
+ library_df.to_csv(out_path("library"), index=False)
+ parameters_df.to_csv(out_path("parameters"), index=False)
+ genotype_theta_df.to_csv(out_path("genotype_theta"), index=False)
+ growth_df.to_csv(out_path("growth"), index=False)
+ print(f"\nWrote: {', '.join(out_path(n) for n in ['library', 'parameters', 'genotype_theta', 'growth'])}")
+
+ if "binding_data" in cf:
+ binding_df = _generate_binding_data(cf["binding_data"], rng, binding_theta_df)
+ binding_df.to_csv(out_path("binding"), index=False)
+ print(f"Wrote: {out_path('binding')}")
+
+ if "presplit_data" in cf:
+ print("\nGenerating presplit data...", flush=True)
+ presplit_df = _generate_presplit_data(combined_sample_df,
+ combined_counts_df,
+ cf, rng)
+ presplit_df.to_csv(out_path("presplit"), index=False)
+ print(f"Wrote: {out_path('presplit')}")
+
+ with open(config_out, "w") as fh:
+ yaml.dump(cf, fh, default_flow_style=False, sort_keys=False)
+ print(f"Wrote: {config_out}")
+
+
+def main():
+ return generalized_main(run_simulation_from_config,
+ manual_arg_types={"seed": int})
diff --git a/src/tfscreen/simulate/selection_experiment.py b/src/tfscreen/simulate/selection_experiment.py
index 77e7f0d4..29eb51a3 100644
--- a/src/tfscreen/simulate/selection_experiment.py
+++ b/src/tfscreen/simulate/selection_experiment.py
@@ -5,6 +5,7 @@
read_dataframe,
read_yaml,
)
+from tfscreen.util.validation import check_unknown_keys
from tfscreen.simulate.growth.transition_linkage import get_transition_model
from tfscreen.util.numerical import (
vstack_padded,
@@ -43,6 +44,30 @@
"min":ma.min,
"max":ma.max,
"sum":ma.sum}
+
+# All recognized top-level keys for a simulate config file.
+SIMULATE_KNOWN_KEYS = frozenset({
+ # Library genetics (passed to LibraryManager)
+ "reading_frame", "first_amplicon_residue", "wt_seq", "degen_sites",
+ "sub_libraries", "library_combos", "spiked_seqs", "expected_5p", "expected_3p",
+ # Phenotype / theta calculation
+ "theta_component", "thermo_data", "theta_priors", "theta_rescale",
+ "theta_sim_priors",
+ # Conditions and growth
+ "condition_blocks", "growth",
+ "dk_geno_hyper_loc", "dk_geno_hyper_scale", "dk_geno_hyper_shift", "dk_geno_zero",
+ "activity_wt", "activity_mut_scale", "activity_component", "activity_priors",
+ # Experimental simulation parameters
+ "transform_sizes", "library_mixture", "lib_assembly_skew_sigma",
+ "transformation_poisson_lambda", "multi_plasmid_combine_fcn", "cfu0",
+ "tube_noise_sigma", "growth_transition",
+ # Data collection
+ "total_num_reads", "prob_index_hop", "seed",
+ # Column selectors (rarely overridden)
+ "condition_selector", "library_selector",
+ # Optional output blocks
+ "binding_data", "presplit_data",
+})
def _check_dict_number(
key: str,
@@ -163,18 +188,16 @@ def _check_cf(
# Load from YAML if a path is provided, otherwise assume it's a dict
cf = read_yaml(cf)
- if "final_cfu_pct_err" not in cf:
- cf["final_cfu_pct_err"] = 0.05
+ check_unknown_keys(cf, SIMULATE_KNOWN_KEYS, label="simulate config")
# --- Validate single numerical values ---
cf = _check_dict_number("prob_index_hop", cf, min_allowed=0, max_allowed=1, allow_none=True)
cf = _check_dict_number("lib_assembly_skew_sigma", cf, min_allowed=0, allow_none=True)
cf = _check_dict_number("transformation_poisson_lambda", cf, min_allowed=0, allow_none=True)
cf = _check_dict_number("tube_noise_sigma", cf, min_allowed=0, allow_none=True)
- cf = _check_dict_number("random_seed", cf, cast_type=int, min_allowed=0, allow_none=True)
+ cf = _check_dict_number("seed", cf, cast_type=int, min_allowed=0, allow_none=True)
cf = _check_dict_number("cfu0", cf, allow_none=False,min_allowed=0)
cf = _check_dict_number("total_num_reads", cf, cast_type=int, min_allowed=0, inclusive_min=False)
- cf = _check_dict_number("final_cfu_pct_err",cf,min_allowed=0,inclusive_min=False,allow_none=False)
# --- Validate nested dictionaries ---
for key in ["transform_sizes", "library_mixture"]:
@@ -981,8 +1004,7 @@ def _simulate_library_group(
multi_plasmid_combine_fcn = cf["multi_plasmid_combine_fcn"]
prob_index_hop = cf["prob_index_hop"]
total_cfu0 = cf["cfu0"]
- final_cfu_pct_err = cf["final_cfu_pct_err"]
-
+
num_genotypes = len(ordered_genotypes)
# -- create output dataframes --
@@ -1084,7 +1106,7 @@ def _simulate_library_group(
# Record cfu/mL over all conditions
sample_df.loc[:,"sample_cfu"] = np.sum(trans_cfu,axis=0)
- sample_df.loc[:,"sample_cfu_std"] = sample_df.loc[:,"sample_cfu"]*final_cfu_pct_err
+ sample_df.loc[:,"sample_cfu_std"] = 0.0
# -- simulate sequencing --
print("--> simulating sequencing",flush=True)
@@ -1196,7 +1218,7 @@ def selection_experiment(
print("Setting up calculation.", flush=True)
# Initialize random number generator
- rng = np.random.default_rng(cf["random_seed"])
+ rng = np.random.default_rng(cf["seed"])
# Expand the library_df so every genotype is seen in every library_origin
# to keep indexing consistent when we mix libraries. The .fillna(0) in this
diff --git a/src/tfscreen/simulate/thermo_to_growth.py b/src/tfscreen/simulate/thermo_to_growth.py
index d13a35a6..ddc47773 100644
--- a/src/tfscreen/simulate/thermo_to_growth.py
+++ b/src/tfscreen/simulate/thermo_to_growth.py
@@ -153,11 +153,74 @@ def _sample_hierarchical_activity(unique_genotypes,
_ACTIVITY_COMPONENTS = {"fixed", "horseshoe_geno", "hierarchical_geno"}
+def _theta_param_to_df(theta_param, unique_genotypes, sim_indices):
+ """
+ Extract per-genotype scalar fields from a ThetaParam pytree.
+
+ Iterates over the dataclass fields of ``theta_param``. Fields whose
+ array has exactly two dimensions (T × G_sim) are treated as
+ per-genotype parameters: columns are selected via ``sim_indices`` and
+ a leading titrant dimension of size 1 is squeezed away; when T > 1 each
+ titrant gets its own column suffixed ``_T{t}``. Fields with any other
+ dimensionality (e.g. population moments ``mu``/``sigma`` with shape
+ T × C × 1) are silently skipped.
+
+ If ``theta_param`` is not a standard or Flax dataclass (e.g. a mock
+ during testing) the function returns a DataFrame containing only the
+ ``genotype`` column, to which ``dk_geno`` and ``activity`` are then
+ added by the caller.
+
+ Parameters
+ ----------
+ theta_param : pytree
+ Returned by ``sample_theta_prior``. Per-genotype fields have shape
+ ``(num_titrant_name, num_genotype_sim)``.
+ unique_genotypes : array-like of str
+ Unique genotype strings in the desired output row order.
+ sim_indices : array-like of int
+ For each entry in ``unique_genotypes``, its column index in the
+ sim_data / ``library_df`` order stored in ``theta_param``.
+
+ Returns
+ -------
+ pandas.DataFrame
+ One row per genotype. Columns are the extracted ``theta_param``
+ field names (with a ``_T{t}`` suffix when ``num_titrant_name > 1``).
+ The ``genotype`` column is always present as the first column.
+ """
+ # Use the class-level __dataclass_fields__ so that MagicMock instances
+ # (which intercept instance-level attribute access) safely return an
+ # empty dict rather than triggering mock attribute creation.
+ field_names = list(
+ getattr(type(theta_param), "__dataclass_fields__", {}).keys()
+ )
+
+ col_data = {}
+ for fname in field_names:
+ try:
+ val = np.array(getattr(theta_param, fname))
+ except Exception:
+ continue
+ if val.ndim != 2:
+ continue # skip mu, sigma (3-D) and any scalars
+ selected = val[:, np.asarray(sim_indices)] # (T, n_unique)
+ if selected.shape[0] == 1:
+ col_data[fname] = selected[0]
+ else:
+ for t in range(selected.shape[0]):
+ col_data[f"{fname}_T{t}"] = selected[t]
+
+ df = pd.DataFrame(col_data)
+ df.insert(0, "genotype", list(unique_genotypes))
+ return df
+
+
def _assign_dk_geno(unique_genotypes,
hyper_loc=-3.5,
hyper_scale=1.0,
hyper_shift=0.02,
- rng: Generator | None = None):
+ rng: Generator | None = None,
+ fixed_value: float | None = None):
"""
Assign a pleiotropic growth-rate effect (dk_geno) to each genotype.
@@ -166,7 +229,13 @@ def _assign_dk_geno(unique_genotypes,
Wild-type receives dk_geno = 0. This matches the prior used in the
hierarchical tfmodel inference component.
+
+ If fixed_value is not None, all genotypes (including wild-type) receive
+ that value and no stochastic sampling is performed.
"""
+ if fixed_value is not None:
+ return pd.Series({g: float(fixed_value) for g in unique_genotypes})
+
if rng is None:
rng = np.random.default_rng()
@@ -225,16 +294,20 @@ def thermo_to_growth(
theta_rng_key,
growth_params: dict,
theta_priors_overrides: Optional[dict] = None,
+ theta_sim_priors_overrides: Optional[dict] = None,
theta_noise_sigma_logit: float = 0.0,
dk_geno_hyper_loc: float = -3.5,
dk_geno_hyper_scale: float = 1.0,
dk_geno_hyper_shift: float = 0.02,
+ dk_geno_zero: bool = False,
activity_wt: float = 1.0,
activity_mut_scale: float = 0.0,
rng: Generator | None = None,
activity_component: str = "fixed",
activity_priors_overrides: Optional[dict] = None,
theta_rescale: str = "passthrough",
+ theta_gc_override: Optional[dict] = None,
+ theta_params_override: Optional[dict] = None,
) -> tuple[pd.DataFrame, pd.DataFrame]:
"""
Generate phenotypes from genotypes via prior-predictive theta sampling.
@@ -255,7 +328,12 @@ def thermo_to_growth(
growth_params : dict
Per-condition growth model parameters keyed by condition string.
theta_priors_overrides : dict or None
- Overrides to the theta component's default hyperparameters.
+ Overrides to the theta component's default ``ModelPriors``
+ hyperparameters (prior-predictive path only; ignored when the
+ component defines ``simulate``).
+ theta_sim_priors_overrides : dict or None
+ Overrides to the theta component's ``SimPriors`` defaults
+ (perturbation path only; ignored for components without ``simulate``).
theta_noise_sigma_logit : float, default 0.0
Standard deviation of additive Normal noise applied on the logit
scale after prior-predictive theta sampling, giving:
@@ -273,6 +351,11 @@ def thermo_to_growth(
dk_geno_hyper_shift : float
Shift applied after exponentiation: dk_geno = hyper_shift - exp(offset).
Controls the maximum possible (beneficial) growth-rate effect.
+ dk_geno_zero : bool, default False
+ When True, assign dk_geno = 0 to every genotype and skip stochastic
+ sampling entirely. The three ``dk_geno_hyper_*`` parameters are
+ ignored. Useful for simulations where no pleiotropic growth effects
+ are desired.
activity_wt : float
TF activity of the wild-type genotype. Used only when
``activity_component == "fixed"``.
@@ -304,17 +387,41 @@ def thermo_to_growth(
the logit scale ``log(θ/(1−θ))``. The ``"theta"`` column in
``phenotype_df`` and ``genotype_theta_df`` always store the
pre-rescale (0–1) value.
+ theta_gc_override : dict[str, np.ndarray] or None, default None
+ Per-genotype theta overrides. Keys are genotype strings; values are
+ 1-D arrays of theta at the growth concentrations (same sorted-unique
+ order as ``sample_df["titrant_conc"]``). Overrides are applied after
+ prior-predictive sampling but before noise and all downstream
+ calculations. Used by ``library_prediction`` to inject
+ stratified theta values for calibration genotypes.
+ theta_params_override : dict[str, dict[str, float]] or None, default None
+ Effective Hill parameters for each overridden genotype. Keys are
+ genotype strings; values are dicts with keys ``theta_low``,
+ ``theta_high``, ``log_hill_K``, ``hill_n``. When provided, the
+ corresponding rows in ``parameters_df`` are updated so that the
+ saved CSV reflects what was actually used in the simulation rather
+ than the prior-predictive sample.
Returns
-------
phenotype_df : pd.DataFrame
Long-form dataframe with one row per (genotype, sample condition).
Columns include ``theta``, ``dk_geno``, ``activity``, ``k_pre``,
- ``k_sel``.
+ ``k_sel``. ``dk_geno`` and ``activity`` are retained here so that
+ ``selection_experiment`` can use them downstream; they are also
+ collected in ``parameters_df`` for file output.
genotype_theta_df : pd.DataFrame
- Wide-form dataframe with one row per genotype and one column per
- unique effector concentration, giving the ground-truth theta value.
+ Long-form dataframe with one row per (genotype, titrant_name,
+ titrant_conc) combination. Columns: ``genotype``, ``titrant_name``,
+ ``titrant_conc``, ``theta``. Sorted by (genotype, titrant_name,
+ titrant_conc); ``genotype`` is an ordered categorical.
Use this to compare simulation inputs against fitted posteriors.
+ parameters_df : pd.DataFrame
+ One row per unique genotype. Always contains ``genotype``,
+ ``dk_geno``, and ``activity``. Also contains any scalar
+ per-genotype fields extracted from the ``theta_param`` pytree
+ (e.g. ``theta_low``, ``theta_high``, ``log_hill_K``, ``hill_n``
+ for ``hill_geno``).
"""
print("Sampling theta from prior... ", end="", flush=True)
@@ -350,10 +457,27 @@ def thermo_to_growth(
sim_data=sim_data,
rng_key=theta_rng_key,
priors_overrides=theta_priors_overrides,
+ sim_priors_overrides=theta_sim_priors_overrides,
)
print("Done.", flush=True)
+ # ── Index lookups (shared by genotype_theta_df, phenotype_df, parameters_df) ─
+ # Computed here (before noise) so that theta_gc_override can be applied
+ # using the same index map. Duplicates in library_df are handled by
+ # mapping each unique genotype to its last occurrence in sim_data order.
+
+ all_genotypes = list(genotypes) # sim_data order (may have duplicates)
+ geno_to_sim_idx = {g: i for i, g in enumerate(all_genotypes)}
+ unique_sim_indices = np.array([geno_to_sim_idx[g] for g in unique_genotypes])
+
+ # ── Inject stratified theta for calibration genotypes ─────────────────────
+ if theta_gc_override:
+ theta_gc = np.array(theta_gc) # ensure a mutable copy
+ for g, theta_row in theta_gc_override.items():
+ if g in geno_to_sim_idx:
+ theta_gc[geno_to_sim_idx[g], :] = theta_row
+
# ── Apply logit-normal noise to theta ─────────────────────────────────────
if theta_noise_sigma_logit > 0.0:
theta_safe = np.clip(theta_gc, 1e-6, 1.0 - 1e-6)
@@ -361,29 +485,45 @@ def thermo_to_growth(
epsilon = rng.normal(0.0, theta_noise_sigma_logit, size=theta_gc.shape)
theta_gc = 1.0 / (1.0 + np.exp(-(logit_theta + epsilon)))
- # ── Build genotype_theta_df (ground-truth parameters for comparison) ─────
- # Wide form: rows = genotypes (in sim_data order), cols = concentrations.
-
- all_genotypes = list(genotypes) # sim_data order
- conc_vals = np.array(sim_data.titrant_conc)
- conc_col_names = [f"theta_at_{c:.6g}mM" for c in conc_vals]
- genotype_theta_df = pd.DataFrame(
- theta_gc,
- index=all_genotypes,
- columns=conc_col_names,
- )
- genotype_theta_df.index.name = "genotype"
- genotype_theta_df = genotype_theta_df.reset_index()
- genotype_theta_df = set_categorical_genotype(genotype_theta_df)
-
# ── Map sample_df concentrations to unique-concentration indices ──────────
# sample_df may have duplicate concentrations across rows; each row maps to
# one column of theta_gc.
- conc_to_col = {float(c): i for i, c in enumerate(conc_vals)}
- sample_conc_idx = sample_df["titrant_conc"].map(conc_to_col).values
+ # Use float64 values from sample_df directly (same source as build_sim_data)
+ # rather than sim_data.titrant_conc, which is float32 from the JAX array and
+ # would cause dict-lookup misses due to float32→float64 precision loss.
+ sorted_concs_f64 = np.sort(sample_df["titrant_conc"].unique()).astype(float)
+ conc_to_col = {c: i for i, c in enumerate(sorted_concs_f64)}
+ sample_conc_idx = sample_df["titrant_conc"].map(conc_to_col).values.astype(int)
+
+ # ── Build genotype_theta_df (ground-truth theta per unique genotype) ──────
+ # Long form: one row per (genotype, titrant_name, titrant_conc).
+ # unique_sim_indices maps each unique genotype to its representative row in
+ # theta_gc, eliminating duplicates from sub-libraries.
+ tn_conc_pairs = (
+ sample_df[["titrant_name", "titrant_conc"]]
+ .drop_duplicates()
+ .sort_values(["titrant_name", "titrant_conc"])
+ .reset_index(drop=True)
+ )
+ pair_col_indices = np.array(
+ [conc_to_col[float(c)] for c in tn_conc_pairs["titrant_conc"]]
+ )
+ # theta_vals shape: (n_unique_geno, n_pairs)
+ theta_vals = theta_gc[unique_sim_indices][:, pair_col_indices]
- # ── Build genotype-index lookup into sim_data order ───────────────────────
- geno_to_sim_idx = {g: i for i, g in enumerate(all_genotypes)}
+ n_geno = len(unique_genotypes)
+ n_pairs = len(tn_conc_pairs)
+ genotype_theta_df = pd.DataFrame({
+ "genotype": np.repeat(list(unique_genotypes), n_pairs),
+ "titrant_name": np.tile(tn_conc_pairs["titrant_name"].values, n_geno),
+ "titrant_conc": np.tile(tn_conc_pairs["titrant_conc"].values.astype(float),
+ n_geno),
+ "theta": theta_vals.ravel(),
+ })
+ genotype_theta_df = set_categorical_genotype(genotype_theta_df)
+ genotype_theta_df = genotype_theta_df.sort_values(
+ ["genotype", "titrant_name", "titrant_conc"]
+ ).reset_index(drop=True)
print("Calculating growth rates and building phenotype dataframe... ",
end="", flush=True)
@@ -396,7 +536,7 @@ def thermo_to_growth(
n_samples = len(sample_df)
# Theta matrix aligned to unique_genotypes order, shape (n_geno, n_samples)
- unique_sim_indices = np.array([geno_to_sim_idx[g] for g in unique_genotypes])
+ # (unique_sim_indices already computed above)
theta_ordered = theta_gc[unique_sim_indices] # (n_geno, C)
theta_for_samples = theta_ordered[:, sample_conc_idx] # (n_geno, n_samples)
@@ -420,6 +560,7 @@ def thermo_to_growth(
genotype_dk_geno_series = _assign_dk_geno(
unique_genotypes, dk_geno_hyper_loc, dk_geno_hyper_scale,
dk_geno_hyper_shift, rng,
+ fixed_value=0.0 if dk_geno_zero else None,
)
phenotype_df["dk_geno"] = phenotype_df["genotype"].map(genotype_dk_geno_series)
@@ -476,6 +617,36 @@ def thermo_to_growth(
final_columns.append("theta")
phenotype_df = phenotype_df.loc[:, final_columns]
+ # ── Build parameters_df (one row per unique genotype) ────────────────────
+ # Extracts per-genotype theta_param fields (e.g. theta_low, log_hill_K)
+ # and pairs them with the already-computed dk_geno and activity values.
+ # phenotype_df keeps dk_geno/activity so selection_experiment still works.
+
+ parameters_df = _theta_param_to_df(theta_param, unique_genotypes,
+ unique_sim_indices)
+ parameters_df["dk_geno"] = genotype_dk_geno_series.reindex(
+ list(unique_genotypes)
+ ).values
+ parameters_df["activity"] = genotype_activity_series.reindex(
+ list(unique_genotypes)
+ ).values
+ _front = ["genotype", "dk_geno", "activity"]
+ _other = [c for c in parameters_df.columns if c not in _front]
+ parameters_df = parameters_df[_front + _other]
+ parameters_df = set_categorical_genotype(parameters_df)
+
+ # Overwrite theta columns with the effective (post-override) Hill parameters
+ # so that parameters_df reflects what was actually simulated, not the
+ # prior-predictive sample that was later replaced.
+ if theta_params_override:
+ for g, effective in theta_params_override.items():
+ mask = parameters_df["genotype"] == g
+ if not mask.any():
+ continue
+ for col, val in effective.items():
+ if col in parameters_df.columns:
+ parameters_df.loc[mask, col] = float(val)
+
print("Done.", flush=True)
- return phenotype_df, genotype_theta_df
+ return phenotype_df, genotype_theta_df, parameters_df
diff --git a/src/tfscreen/tfmodel/analysis/error_calibration.py b/src/tfscreen/tfmodel/analysis/error_calibration.py
new file mode 100644
index 00000000..ea6f3cf3
--- /dev/null
+++ b/src/tfscreen/tfmodel/analysis/error_calibration.py
@@ -0,0 +1,644 @@
+"""
+Error calibration diagnostics for posterior distributions.
+
+Theory
+------
+A well-calibrated posterior assigns the correct probability to every interval:
+if the model says the 90 % credible interval covers the true value 90 % of the
+time, it should. The central diagnostic tool is the **Probability Integral
+Transform (PIT)**: for a single parameter with posterior CDF F and true value
+θ*, the PIT is F(θ*). For a perfectly calibrated model, PIT values drawn
+across many parameters or many independent datasets are Uniform(0, 1).
+
+This module supports two calibration workflows that share the same core:
+
+1. **Single-run calibration** (``calibration_summary``): one inference run on
+ simulated data with known ground truth. True values come from the fixed
+ simulation; posterior uncertainty comes from the marginal quantile CSVs
+ written by ``tfs-param-quantiles``. The PIT for each genotype/parameter
+ is interpolated from the stored quantiles.
+
+2. **Multi-run SBC** (``summarize_sbc``, Simulation-Based Calibration per
+ Talts et al. 2018): many independent inference runs, each with a dataset
+ simulated from a fresh prior draw. The PIT for each parameter element is
+ computed from the full posterior sample array. Pooling PIT values across
+ runs and testing uniformity validates that the inference algorithm correctly
+ implements the generative model's prior.
+
+Functional layers
+-----------------
+- **Core** : ``pit_from_samples``, ``pit_from_quantiles``
+- **Stats** : ``calibration_curve``, ``pit_uniformity_test``
+- **Plots** : ``plot_pit_histogram``, ``plot_calibration_curve``
+- **Single-run** : ``calibration_summary``
+- **SBC** : ``compute_sbc_ranks``, ``summarize_sbc``
+
+References
+----------
+Talts et al. (2018) "Validating Bayesian Inference Algorithms with
+Simulation-Based Calibration". arXiv:1804.06788.
+"""
+
+import glob
+import os
+import warnings
+from typing import Dict, List, Optional, Tuple
+
+import h5py
+import numpy as np
+import pandas as pd
+from matplotlib import pyplot as plt
+from scipy.stats import ks_1samp, uniform
+
+
+# ============================================================
+# Part 1 — Core PIT computation
+# ============================================================
+
+def pit_from_samples(true_vals, posterior_samples):
+ """Compute PIT values from raw posterior samples.
+
+ For each element the PIT is the fraction of posterior samples strictly
+ less than the true value — equivalently, the posterior rank of the true
+ value. For a perfectly calibrated model these values are Uniform(0, 1).
+
+ Parameters
+ ----------
+ true_vals : array-like, shape (*D)
+ True (ground-truth) parameter values. NaN entries propagate as NaN
+ in the output.
+ posterior_samples : array-like, shape (S, *D)
+ S posterior samples for each element of *D*.
+
+ Returns
+ -------
+ ndarray, shape (*D)
+ PIT values in [0, 1], NaN where ``true_vals`` is NaN.
+ """
+ true_vals = np.asarray(true_vals, dtype=float)
+ posterior_samples = np.asarray(posterior_samples, dtype=float)
+
+ pit = np.mean(posterior_samples < true_vals, axis=0).astype(float)
+
+ nan_mask = np.isnan(true_vals)
+ if np.any(nan_mask):
+ pit = np.where(nan_mask, np.nan, pit)
+
+ return pit
+
+
+def pit_from_quantiles(true_vals, quantile_matrix, quantile_levels):
+ """Interpolate PIT values from stored posterior quantiles.
+
+ For each observation i the PIT is estimated by linear interpolation of the
+ posterior CDF at ``true_vals[i]``. Values below the lowest stored
+ quantile receive PIT = 0; values above the highest receive PIT = 1.
+
+ Parameters
+ ----------
+ true_vals : array-like, shape (N,)
+ True (ground-truth) values, one per observation.
+ quantile_matrix : array-like, shape (N, Q)
+ Posterior quantile values; row i contains the Q quantile estimates for
+ observation i. Columns must be in ascending order of ``quantile_levels``.
+ quantile_levels : array-like, shape (Q,)
+ Probability levels corresponding to the columns of ``quantile_matrix``,
+ e.g. ``[0.025, 0.25, 0.50, 0.75, 0.975]``. Must be sorted ascending.
+
+ Returns
+ -------
+ ndarray, shape (N,)
+ Interpolated PIT values in [0, 1]. NaN where ``true_vals[i]`` is NaN
+ or any quantile for row i is NaN.
+
+ Notes
+ -----
+ Accuracy depends on the density of stored quantile levels; coarse grids
+ (e.g. only 2.5 % and 97.5 %) produce step-function PIT estimates.
+ """
+ true_vals = np.asarray(true_vals, dtype=float)
+ quantile_matrix = np.asarray(quantile_matrix, dtype=float)
+ quantile_levels = np.asarray(quantile_levels, dtype=float)
+
+ N = len(true_vals)
+ pit = np.full(N, np.nan)
+
+ for i in range(N):
+ tv = true_vals[i]
+ qv = quantile_matrix[i]
+ if np.isnan(tv) or np.any(np.isnan(qv)):
+ continue
+ # np.interp: xp must be increasing; quantile values are non-decreasing
+ pit[i] = np.interp(tv, qv, quantile_levels, left=0.0, right=1.0)
+
+ return pit
+
+
+# ============================================================
+# Part 2 — Calibration statistics
+# ============================================================
+
+_DEFAULT_CALIBRATION_LEVELS = np.array([0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99])
+
+
+def calibration_curve(pit_vals, levels=None):
+ """Compute empirical coverage at multiple nominal credible-interval levels.
+
+ For each nominal level α the empirical coverage is the fraction of PIT
+ values that fall in the symmetric interval [(1−α)/2, (1+α)/2]. For a
+ perfectly calibrated model the empirical coverage equals α at every level,
+ so plotting empirical vs. nominal should lie on the diagonal.
+
+ Parameters
+ ----------
+ pit_vals : array-like
+ PIT values in [0, 1]. Non-finite values are ignored.
+ levels : array-like, optional
+ Nominal coverage levels to evaluate, e.g. ``[0.5, 0.9, 0.95]``.
+ Defaults to ``[0.50, 0.60, 0.70, 0.80, 0.90, 0.95, 0.99]``.
+
+ Returns
+ -------
+ dict
+ Mapping ``{nominal_level (float): empirical_coverage (float)}``.
+ """
+ if levels is None:
+ levels = _DEFAULT_CALIBRATION_LEVELS
+ pit = np.asarray(pit_vals, dtype=float)
+ pit = pit[np.isfinite(pit)]
+
+ result = {}
+ for alpha in levels:
+ lower = (1.0 - float(alpha)) / 2.0
+ upper = 1.0 - lower
+ empirical = float(np.mean((pit >= lower) & (pit <= upper)))
+ result[float(alpha)] = empirical
+ return result
+
+
+def pit_uniformity_test(pit_vals):
+ """One-sample KS test of PIT values against Uniform(0, 1).
+
+ Parameters
+ ----------
+ pit_vals : array-like
+ PIT values. Non-finite values are ignored.
+
+ Returns
+ -------
+ dict
+ ``ks_stat`` (float), ``ks_pval`` (float), ``mean_pit`` (float),
+ ``n_vals`` (int). All numeric fields are NaN when no finite values
+ are present.
+ """
+ pit = np.asarray(pit_vals, dtype=float)
+ pit = pit[np.isfinite(pit)]
+ n = int(len(pit))
+ if n == 0:
+ return {"ks_stat": np.nan, "ks_pval": np.nan, "mean_pit": np.nan, "n_vals": 0}
+ mean_pit = float(np.mean(pit))
+ ks_stat, ks_pval = ks_1samp(pit, uniform(0, 1).cdf)
+ return {
+ "ks_stat": float(ks_stat),
+ "ks_pval": float(ks_pval),
+ "mean_pit": mean_pit,
+ "n_vals": n,
+ }
+
+
+# ============================================================
+# Part 3 — Plotting
+# ============================================================
+
+def plot_pit_histogram(pit_vals, ax=None, title=None, n_bins=10):
+ """Plot a PIT histogram with a uniform-density reference line.
+
+ Parameters
+ ----------
+ pit_vals : array-like
+ PIT values. Non-finite values are silently dropped.
+ ax : matplotlib.axes.Axes, optional
+ Axes to draw on. A new figure is created when *None*.
+ title : str, optional
+ Axes title.
+ n_bins : int
+ Number of histogram bins over [0, 1] (default 10).
+
+ Returns
+ -------
+ matplotlib.axes.Axes
+ """
+ pit = np.asarray(pit_vals, dtype=float)
+ pit = pit[np.isfinite(pit)]
+
+ if ax is None:
+ _, ax = plt.subplots(1, 1, figsize=(5, 4))
+
+ ax.hist(pit, bins=n_bins, range=(0, 1), density=True,
+ color="steelblue", alpha=0.75, edgecolor="white")
+ ax.axhline(1.0, color="firebrick", lw=1.5, ls="--", label="Uniform")
+ ax.set_xlim(0, 1)
+ ax.set_ylim(0, n_bins)
+ ax.set_xlabel("PIT value")
+ ax.set_ylabel("Density")
+ if title is not None:
+ ax.set_title(title)
+ ax.legend(fontsize=8)
+ return ax
+
+
+def plot_calibration_curve(calibration_dict, ax=None, label=None):
+ """Plot empirical vs. nominal coverage (calibration curve).
+
+ A perfectly calibrated model lies on the diagonal. Points above the
+ diagonal indicate underconfidence (CIs too wide); points below indicate
+ overconfidence (CIs too narrow).
+
+ Parameters
+ ----------
+ calibration_dict : dict
+ ``{nominal_level: empirical_coverage}`` as returned by
+ ``calibration_curve``.
+ ax : matplotlib.axes.Axes, optional
+ Axes to draw on. A new figure is created when *None*.
+ label : str, optional
+ Legend label for the model curve.
+
+ Returns
+ -------
+ matplotlib.axes.Axes
+ """
+ if ax is None:
+ _, ax = plt.subplots(1, 1, figsize=(5, 5))
+
+ if not calibration_dict:
+ ax.set_xlabel("Nominal coverage")
+ ax.set_ylabel("Empirical coverage")
+ return ax
+
+ nominals = np.array(sorted(calibration_dict.keys()))
+ empiricals = np.array([calibration_dict[k] for k in nominals])
+
+ ax.plot([-0.05, 1.05], [-0.05, 1.05], "k--", lw=1, label="Ideal")
+ ax.plot(nominals, empiricals, "o-", color="steelblue",
+ label=label or "Model")
+ ax.set_xlim(-0.05, 1.05)
+ ax.set_ylim(-0.05, 1.05)
+ ax.set_xlabel("Nominal coverage")
+ ax.set_ylabel("Empirical coverage")
+ ax.legend(fontsize=8)
+ ax.set_aspect("equal")
+ return ax
+
+
+# ============================================================
+# Part 4 — Single-run calibration summary
+# ============================================================
+
+def calibration_summary(true_vals, quantile_matrix, quantile_levels,
+ out_prefix, label=""):
+ """Compute and write calibration diagnostics for a single inference run.
+
+ Computes PIT values via interpolation from stored quantiles, then writes
+ a PIT CSV, a calibration-curve CSV, and a two-panel PDF.
+
+ **CSV files are the authoritative outputs.** The PDF is a human-readable
+ summary; the CSVs are the record of truth for downstream analysis.
+
+ Output files
+ ------------
+ ``{out_prefix}_pit.csv``
+ Columns: ``true_val``, ``pit``.
+ ``{out_prefix}_calibration_curve.csv``
+ Columns: ``nominal``, ``empirical``.
+ ``{out_prefix}_calibration.pdf``
+ Two panels: PIT histogram (left) and calibration curve (right).
+
+ Parameters
+ ----------
+ true_vals : array-like, shape (N,)
+ Known ground-truth values from simulation.
+ quantile_matrix : array-like, shape (N, Q)
+ Posterior quantile values; row i is quantile estimates for observation
+ i at levels ``quantile_levels``.
+ quantile_levels : array-like, shape (Q,)
+ Probability levels (sorted ascending) corresponding to columns of
+ ``quantile_matrix``.
+ out_prefix : str
+ File-system prefix for all output files.
+ label : str, optional
+ Human-readable label used in plot titles.
+
+ Returns
+ -------
+ dict
+ ``{"n_vals": int, "ks_stat": float, "ks_pval": float,
+ "mean_pit": float, "calibration_curve": {nominal: empirical, ...}}``
+ """
+ os.makedirs(os.path.dirname(os.path.abspath(out_prefix)), exist_ok=True)
+
+ true_vals = np.asarray(true_vals, dtype=float)
+ pit_vals = pit_from_quantiles(true_vals, quantile_matrix, quantile_levels)
+ cal_curve = calibration_curve(pit_vals)
+ uniformity = pit_uniformity_test(pit_vals)
+
+ # --- PIT CSV ---
+ pit_csv = f"{out_prefix}_pit.csv"
+ try:
+ pd.DataFrame({"true_val": true_vals, "pit": pit_vals}).to_csv(pit_csv, index=False)
+ print(f"Wrote PIT values to {pit_csv}")
+ except Exception as exc:
+ warnings.warn(f"Could not write PIT CSV to {pit_csv}: {exc}")
+
+ # --- Calibration-curve CSV ---
+ cal_csv = f"{out_prefix}_calibration_curve.csv"
+ try:
+ pd.DataFrame({
+ "nominal": list(cal_curve.keys()),
+ "empirical": list(cal_curve.values()),
+ }).to_csv(cal_csv, index=False)
+ print(f"Wrote calibration curve to {cal_csv}")
+ except Exception as exc:
+ warnings.warn(f"Could not write calibration curve CSV to {cal_csv}: {exc}")
+
+ # --- Two-panel PDF ---
+ pdf_path = f"{out_prefix}_calibration.pdf"
+ try:
+ fig, axes = plt.subplots(1, 2, figsize=(10, 5))
+ hist_title = f"PIT histogram — {label}" if label else "PIT histogram"
+ curve_title = f"Calibration curve — {label}" if label else "Calibration curve"
+ plot_pit_histogram(pit_vals, ax=axes[0], title=hist_title)
+ plot_calibration_curve(cal_curve, ax=axes[1])
+ axes[1].set_title(curve_title)
+ fig.tight_layout()
+ fig.savefig(pdf_path, format="pdf", bbox_inches="tight")
+ plt.close(fig)
+ print(f"Wrote calibration plot to {pdf_path}")
+ except Exception as exc:
+ warnings.warn(f"Could not write calibration PDF to {pdf_path}: {exc}")
+
+ return {
+ "n_vals": uniformity["n_vals"],
+ "ks_stat": uniformity["ks_stat"],
+ "ks_pval": uniformity["ks_pval"],
+ "mean_pit": uniformity["mean_pit"],
+ "calibration_curve": cal_curve,
+ }
+
+
+# ============================================================
+# Part 5 — Multi-run SBC (Simulation-Based Calibration)
+# ============================================================
+
+def _load_h5_params(path: str) -> Dict[str, np.ndarray]:
+ """Load all dataset arrays from an HDF5 file."""
+ out = {}
+ with h5py.File(path, "r") as hf:
+ for key in hf.keys():
+ out[key] = hf[key][:]
+ return out
+
+
+def compute_sbc_ranks(ground_truth_path: str,
+ posterior_path: str) -> Dict[str, np.ndarray]:
+ """Compute the posterior rank (PIT) of each ground-truth parameter value.
+
+ For each parameter present in both files the rank is the fraction of
+ posterior samples strictly less than the ground-truth value::
+
+ rank[i] = mean(posterior_samples[:, i] < gt_value[i]) ∈ [0, 1]
+
+ Parameters
+ ----------
+ ground_truth_path : str
+ HDF5 file produced by ``tfs-sample-prior``. Each dataset has shape
+ ``(1, *param_shape)`` (num_draws=1).
+ posterior_path : str
+ HDF5 file produced by ``tfs-sample-posterior``. Each dataset has
+ shape ``(S, *param_shape)`` with S posterior samples.
+
+ Returns
+ -------
+ dict
+ Mapping from parameter name to a 1-D rank array (one value per scalar
+ element of the parameter). Only parameters present in both files are
+ included.
+ """
+ gt_params = _load_h5_params(ground_truth_path)
+ post_params = _load_h5_params(posterior_path)
+
+ ranks = {}
+ for name, gt_val in gt_params.items():
+ if name not in post_params:
+ continue
+
+ gt = gt_val[0] if gt_val.ndim >= 1 else gt_val
+ post = post_params[name]
+
+ if post.ndim < 1 or post.shape[0] == 0:
+ continue
+
+ try:
+ gt_f = np.asarray(gt, dtype=float)
+ post_f = np.asarray(post, dtype=float)
+ except (TypeError, ValueError):
+ continue
+
+ ranks[name] = pit_from_samples(gt_f, post_f).ravel()
+
+ return ranks
+
+
+def _find_pairs(sbc_dir: str) -> List[Tuple[str, str, Optional[str]]]:
+ """Find (run_id, ground_truth_path, posterior_path_or_None) triples.
+
+ Scans *sbc_dir* for files matching ``*_ground_truth.h5``. For each,
+ looks for a corresponding ``*_posterior.h5`` with the same prefix
+ (replacing ``_ground_truth`` with ``_posterior``).
+ """
+ gt_files = sorted(glob.glob(os.path.join(sbc_dir, "*_ground_truth.h5")))
+ pairs = []
+ for gt_path in gt_files:
+ basename = os.path.basename(gt_path)
+ run_id = basename[: -len("_ground_truth.h5")]
+ post_path = os.path.join(sbc_dir, f"{run_id}_posterior.h5")
+ pairs.append((run_id, gt_path,
+ post_path if os.path.exists(post_path) else None))
+ return pairs
+
+
+_SBC_HIST_BINS = 10
+
+
+def _plot_sbc_rank_histograms(all_ranks: Dict[str, List[np.ndarray]],
+ out_prefix: str) -> None:
+ """Write a PDF with one PIT-histogram panel per parameter."""
+ params = sorted(all_ranks)
+ n = len(params)
+ if n == 0:
+ return
+
+ ncols = min(4, n)
+ nrows = (n + ncols - 1) // ncols
+ fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3 * nrows),
+ squeeze=False)
+
+ for ax_idx, param in enumerate(params):
+ ax = axes[ax_idx // ncols][ax_idx % ncols]
+ pooled = np.concatenate(all_ranks[param])
+ pooled = pooled[np.isfinite(pooled)]
+
+ plot_pit_histogram(pooled, ax=ax, n_bins=_SBC_HIST_BINS)
+
+ title = param if len(param) <= 30 else f"…{param[-27:]}"
+ _, ks_pval = ks_1samp(pooled, uniform(0, 1).cdf) if len(pooled) > 0 else (np.nan, np.nan)
+ ax.set_title(f"{title}\n(n={len(pooled)}, p={ks_pval:.2g})", fontsize=9)
+
+ for ax_idx in range(n, nrows * ncols):
+ axes[ax_idx // ncols][ax_idx % ncols].set_visible(False)
+
+ fig.tight_layout()
+ pdf_path = f"{out_prefix}_rank_hist.pdf"
+ try:
+ fig.savefig(pdf_path)
+ print(f"Wrote rank histograms → {pdf_path}", flush=True)
+ except Exception as exc:
+ warnings.warn(f"Could not write rank histogram PDF: {exc}")
+ finally:
+ plt.close(fig)
+
+
+def summarize_sbc(sbc_dir: str, out_prefix: Optional[str] = None) -> pd.DataFrame:
+ """Scan *sbc_dir* for SBC run pairs and compute rank-uniformity statistics.
+
+ For each ``*_ground_truth.h5`` / ``*_posterior.h5`` pair found, compute
+ the posterior rank (PIT) of every ground-truth parameter value. Pool
+ ranks across runs and write three output files:
+
+ ``{out_prefix}_sbc_summary.csv``
+ One row per parameter with mean rank, KS statistic, and p-value.
+ Mean rank near 0.5 and high p-value indicate good calibration.
+
+ ``{out_prefix}_sbc_ranks.csv``
+ Long-form table of every individual rank value.
+ Columns: ``run_id``, ``param``, ``element_idx``, ``rank``.
+
+ ``{out_prefix}_rank_hist.pdf``
+ One-panel-per-parameter PIT histogram. A flat histogram indicates
+ uniform ranks (good calibration).
+
+ Parameters
+ ----------
+ sbc_dir : str
+ Directory containing ``*_ground_truth.h5`` and ``*_posterior.h5``
+ files produced by ``tfs-sample-prior`` and ``tfs-sample-posterior``.
+ out_prefix : str, optional
+ Prefix for all output files. Defaults to ``{sbc_dir}/sbc``.
+
+ Returns
+ -------
+ pd.DataFrame
+ Per-parameter summary table (same content as the CSV).
+ """
+ sbc_dir = os.path.abspath(sbc_dir)
+ if not os.path.isdir(sbc_dir):
+ raise FileNotFoundError(f"SBC directory not found: {sbc_dir}")
+
+ if out_prefix is None:
+ out_prefix = os.path.join(sbc_dir, "sbc")
+
+ pairs = _find_pairs(sbc_dir)
+ if not pairs:
+ print(f"No *_ground_truth.h5 files found in {sbc_dir}/", flush=True)
+ return pd.DataFrame()
+
+ n_total = len(pairs)
+ n_posterior = sum(1 for _, _, p in pairs if p is not None)
+ print(
+ f"Found {n_total} ground-truth file(s), "
+ f"{n_posterior} with matching posterior.",
+ flush=True,
+ )
+ if n_posterior == 0:
+ warnings.warn(
+ "No posterior files found. Run tfs-sample-posterior for each "
+ "synthetic dataset before calling tfs-summarize-sbc."
+ )
+ return pd.DataFrame()
+
+ all_ranks: Dict[str, List[np.ndarray]] = {}
+ rank_rows = []
+
+ for run_id, gt_path, post_path in pairs:
+ if post_path is None:
+ continue
+ try:
+ run_ranks = compute_sbc_ranks(gt_path, post_path)
+ except Exception as exc:
+ warnings.warn(f"Skipping run '{run_id}': {exc}")
+ continue
+
+ for param, rank_arr in run_ranks.items():
+ all_ranks.setdefault(param, []).append(rank_arr)
+ for idx, r in enumerate(rank_arr):
+ rank_rows.append({
+ "run_id": run_id,
+ "param": param,
+ "element_idx": idx,
+ "rank": float(r),
+ })
+
+ if not all_ranks:
+ warnings.warn("No matching parameters found across ground-truth and posterior files.")
+ return pd.DataFrame()
+
+ summary_rows = []
+ for param in sorted(all_ranks):
+ pooled = np.concatenate(all_ranks[param])
+ uniformity = pit_uniformity_test(pooled)
+ n_runs = len(all_ranks[param])
+ summary_rows.append({
+ "param": param,
+ "n_runs": n_runs,
+ "n_ranks": uniformity["n_vals"],
+ "mean_rank": uniformity["mean_pit"],
+ "expected_rank": 0.5,
+ "ks_stat": uniformity["ks_stat"],
+ "ks_pval": uniformity["ks_pval"],
+ })
+
+ summary_df = pd.DataFrame(summary_rows)
+
+ summary_csv = f"{out_prefix}_sbc_summary.csv"
+ summary_df.to_csv(summary_csv, index=False)
+ print(f"Wrote per-parameter summary → {summary_csv}", flush=True)
+
+ if rank_rows:
+ ranks_csv = f"{out_prefix}_sbc_ranks.csv"
+ ranks_df = pd.DataFrame(rank_rows)
+ ranks_df.to_csv(ranks_csv, index=False)
+ print(f"Wrote all rank values → {ranks_csv}", flush=True)
+
+ _plot_sbc_rank_histograms(all_ranks, out_prefix)
+
+ print(
+ f"\nSummary: {len(summary_df)} parameter(s) across "
+ f"{n_posterior} run(s).",
+ flush=True,
+ )
+ low_pval = summary_df[summary_df["ks_pval"] < 0.05] if len(summary_df) else pd.DataFrame()
+ if len(low_pval):
+ print(
+ f" {len(low_pval)} parameter(s) with KS p-value < 0.05 "
+ f"(possible miscalibration):",
+ flush=True,
+ )
+ for _, row in low_pval.iterrows():
+ print(
+ f" {row['param']:40s} KS={row['ks_stat']:.3f} p={row['ks_pval']:.3g}",
+ flush=True,
+ )
+ else:
+ print(" All parameters pass KS uniformity test (p ≥ 0.05).", flush=True)
+
+ return summary_df
diff --git a/src/tfscreen/tfmodel/analysis/extraction.py b/src/tfscreen/tfmodel/analysis/extraction.py
index d0be96a1..220dcf08 100644
--- a/src/tfscreen/tfmodel/analysis/extraction.py
+++ b/src/tfscreen/tfmodel/analysis/extraction.py
@@ -100,7 +100,7 @@ def _extract_param_est(input_df,
return out_dfs
-def extract_parameters(model, posteriors, q_to_get=None):
+def extract_parameters(orchestrator, posteriors, q_to_get=None):
"""
Extract parameter quantiles from posterior samples.
@@ -110,12 +110,12 @@ def extract_parameters(model, posteriors, q_to_get=None):
Parameters
----------
- model : ModelOrchestrator
+ orchestrator : ModelOrchestrator
The model instance to extract parameters from.
posteriors : dict or str
- Assumes this is a dictionary of posteriors keying parameters to
- numpy arrays, a numpy.lib.npyio.NpzFile object, or a path to a
- .npz or .h5/.hdf5 file containing posterior samples for model
+ Assumes this is a dictionary of posteriors keying parameters to
+ numpy arrays, a numpy.lib.npyio.NpzFile object, or a path to a
+ .npz or .h5/.hdf5 file containing posterior samples for model
parameters.
q_to_get : dict, optional
Dictionary mapping output column names to quantile values (between 0 and 1)
@@ -138,20 +138,20 @@ def extract_parameters(model, posteriors, q_to_get=None):
q_to_get, param_posteriors = load_posteriors(posteriors, q_to_get)
ctx = ExtractionContext(
- growth_tm=model.training_tm,
- mut_labels=model.mut_labels,
- pair_labels=model.pair_labels,
- growth_shares_replicates=model._growth_shares_replicates,
+ growth_tm=orchestrator.training_tm,
+ mut_labels=orchestrator.mut_labels,
+ pair_labels=orchestrator.pair_labels,
+ growth_shares_replicates=orchestrator._growth_shares_replicates,
)
component_selections = [
- ("condition_growth", model._condition_growth),
- ("growth_transition", model._growth_transition),
+ ("condition_growth", orchestrator._condition_growth),
+ ("growth_transition", orchestrator._growth_transition),
("ln_cfu0", "hierarchical"),
- ("dk_geno", model._dk_geno),
- ("activity", model._activity),
- ("theta", model._theta),
- ("transformation", model._transformation),
+ ("dk_geno", orchestrator._dk_geno),
+ ("activity", orchestrator._activity),
+ ("theta", orchestrator._theta),
+ ("transformation", orchestrator._transformation),
]
extract = []
@@ -179,7 +179,7 @@ def extract_parameters(model, posteriors, q_to_get=None):
return params
-def extract_theta_curves(model, posteriors, q_to_get=None, manual_titrant_df=None,
+def extract_theta_curves(orchestrator, posteriors, q_to_get=None, manual_titrant_df=None,
num_samples=100):
"""
Extract theta curves by sampling from the joint posterior distribution.
@@ -190,7 +190,7 @@ def extract_theta_curves(model, posteriors, q_to_get=None, manual_titrant_df=Non
Parameters
----------
- model : ModelOrchestrator
+ orchestrator : ModelOrchestrator
The model instance to extract parameters from.
posteriors : dict or str
Assumes this is a dictionary of posteriors keying parameters to
@@ -206,7 +206,7 @@ def extract_theta_curves(model, posteriors, q_to_get=None, manual_titrant_df=Non
at which to calculate theta. If provided, it overrides the default
calculation at the concentrations present in the input data.
If 'genotype' is present, it will be used; otherwise, the method
- will calculate theta for all genotypes in the model.
+ will calculate theta for all genotypes in the orchestrator.
num_samples : int or None, optional
Randomly select this many joint posterior samples and return them as
columns ``sample_0``, ``sample_1``, ... alongside the quantile
@@ -228,18 +228,18 @@ def extract_theta_curves(model, posteriors, q_to_get=None, manual_titrant_df=Non
If `q_to_get` is not a dictionary.
If `manual_titrant_df` is missing required columns.
"""
- module = model_registry.get("theta", {}).get(model._theta)
+ module = model_registry.get("theta", {}).get(orchestrator._theta)
if module is None or not (hasattr(module, "build_calc_df")
and hasattr(module, "compute_theta_samples")):
raise ValueError(
f"extract_theta_curves requires the theta component to implement "
f"build_calc_df and compute_theta_samples. "
- f"'{model._theta}' does not support this interface."
+ f"'{orchestrator._theta}' does not support this interface."
)
q_to_get, param_posteriors = load_posteriors(posteriors, q_to_get)
- calc_df, internal_cols, extra_kwargs = module.build_calc_df(model, manual_titrant_df)
+ calc_df, internal_cols, extra_kwargs = module.build_calc_df(orchestrator, manual_titrant_df)
theta_samples = module.compute_theta_samples(calc_df, param_posteriors, **extra_kwargs)
for q_name, q_val in q_to_get.items():
@@ -257,7 +257,7 @@ def extract_theta_curves(model, posteriors, q_to_get=None, manual_titrant_df=Non
return calc_df.drop(columns=internal_cols)
-def extract_theta_unmeasured(model, posteriors, target_genotypes,
+def extract_theta_unmeasured(orchestrator, posteriors, target_genotypes,
manual_titrant_df, q_to_get=None,
genotype_batch_size=2000):
"""
@@ -273,7 +273,7 @@ def extract_theta_unmeasured(model, posteriors, target_genotypes,
Parameters
----------
- model : ModelOrchestrator
+ orchestrator : ModelOrchestrator
Fitted model instance (must carry ``mut_labels``, ``pair_labels``,
``growth_tm``, and ``priors``).
posteriors : dict or str
@@ -284,7 +284,7 @@ def extract_theta_unmeasured(model, posteriors, target_genotypes,
manual_titrant_df : pd.DataFrame
Must have ``'titrant_name'`` and ``'titrant_conc'`` columns specifying
the concentration grid. All ``titrant_name`` values must be present in
- the model.
+ the orchestrator.
q_to_get : dict, optional
Quantiles to extract. Defaults to the standard set used by other
extraction functions.
@@ -305,20 +305,20 @@ def extract_theta_unmeasured(model, posteriors, target_genotypes,
ValueError
If the active theta component does not implement ``predict_unmeasured``.
"""
- module = model_registry.get("theta", {}).get(model._theta)
+ module = model_registry.get("theta", {}).get(orchestrator._theta)
if module is None or not hasattr(module, "predict_unmeasured"):
raise ValueError(
f"extract_theta_unmeasured requires the theta component to implement "
- f"predict_unmeasured. '{model._theta}' does not support this interface."
+ f"predict_unmeasured. '{orchestrator._theta}' does not support this interface."
)
q_to_get, param_posteriors = load_posteriors(posteriors, q_to_get)
- titrant_dim = model.training_tm.tensor_dim_names.index("titrant_name")
- titrant_names = list(model.training_tm.tensor_dim_labels[titrant_dim])
+ titrant_dim = orchestrator.training_tm.tensor_dim_names.index("titrant_name")
+ titrant_names = list(orchestrator.training_tm.tensor_dim_labels[titrant_dim])
extra_kwargs = {}
- theta_priors = model.priors.theta
+ theta_priors = orchestrator.priors.theta
if hasattr(theta_priors, "theta_tf_total_M"):
extra_kwargs["tf_total"] = float(theta_priors.theta_tf_total_M)
if hasattr(theta_priors, "theta_op_total_M"):
@@ -334,8 +334,8 @@ def extract_theta_unmeasured(model, posteriors, target_genotypes,
target_genotypes=target_genotypes,
titrant_names=titrant_names,
manual_titrant_df=manual_titrant_df,
- mut_labels=model.mut_labels,
- pair_labels=model.pair_labels,
+ mut_labels=orchestrator.mut_labels,
+ pair_labels=orchestrator.pair_labels,
param_posteriors=param_posteriors,
q_to_get=q_to_get,
**extra_kwargs,
@@ -355,8 +355,8 @@ def extract_theta_unmeasured(model, posteriors, target_genotypes,
target_genotypes=batch,
titrant_names=titrant_names,
manual_titrant_df=manual_titrant_df,
- mut_labels=model.mut_labels,
- pair_labels=model.pair_labels,
+ mut_labels=orchestrator.mut_labels,
+ pair_labels=orchestrator.pair_labels,
param_posteriors=param_posteriors,
q_to_get=q_to_get,
**extra_kwargs,
@@ -366,7 +366,7 @@ def extract_theta_unmeasured(model, posteriors, target_genotypes,
return pd.concat(result_dfs, ignore_index=True)
-def extract_growth_predictions(model,
+def extract_growth_predictions(orchestrator,
posteriors,
q_to_get=None,
num_samples=100):
@@ -374,11 +374,11 @@ def extract_growth_predictions(model,
Extract predicted ln_cfu values matching the input growth data.
This method pulls the 'growth_pred' values from the posterior samples
- and maps them back to the original rows in `model.growth_df`.
+ and maps them back to the original rows in `orchestrator.growth_df`.
Parameters
----------
- model : ModelOrchestrator
+ orchestrator : ModelOrchestrator
The model instance to extract parameters from.
posteriors : dict or str
Assumes this is a dictionary of posteriors keying parameters to
@@ -401,7 +401,7 @@ def extract_growth_predictions(model,
Returns
-------
pd.DataFrame
- A copy of `model.growth_df` with new columns for the requested
+ A copy of `orchestrator.growth_df` with new columns for the requested
quantiles of 'ln_cfu_pred' and (unless ``num_samples`` is ``None``)
``sample_0`` … ``sample_{num_samples-1}``.
@@ -429,7 +429,7 @@ def extract_growth_predictions(model,
group_cols = ["replicate_idx", "time_idx", "condition_pre_idx",
"condition_sel_idx", "titrant_name_idx", "titrant_conc_idx"]
- out_df = model.growth_df.copy()
+ out_df = orchestrator.growth_df.copy()
out_df = out_df.sort_values(by=group_cols + ["genotype_idx"])
keep_columns = ["replicate", "genotype",
diff --git a/src/tfscreen/tfmodel/analysis/prediction.py b/src/tfscreen/tfmodel/analysis/prediction.py
index 79294f87..090b6093 100644
--- a/src/tfscreen/tfmodel/analysis/prediction.py
+++ b/src/tfscreen/tfmodel/analysis/prediction.py
@@ -3,45 +3,109 @@
import itertools
from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
from tfscreen.tfmodel.inference.posteriors import load_posteriors, get_posterior_samples
+from tfscreen.tfmodel.tensors.batch import get_batch as _get_batch
import jax
from jax import numpy as jnp
from numpyro.handlers import seed, trace
from numpyro.infer import Predictive
-def copy_model_class(model_class,
- t_pre=None,
- t_sel=None,
- titrant_conc=None):
+def _align_site_to_tm_dims(spatial_shape, tm_sizes):
+ """
+ Map each axis of a predicted site's spatial_shape to the correct
+ TensorManager dimension index.
+
+ Right-to-left alignment is tried first. It is correct when the
+ non-broadcast site dimensions are tail-aligned with the TM ordering —
+ e.g. growth_pred (7-D, full TM rank) or theta_growth_pred (7-D with
+ broadcast size-1 dims on the left).
+
+ For sites whose non-broadcast dims are NOT tail-aligned — notably ln_cfu0
+ with shape (n_rep, n_cp, n_geno) where n_rep and n_cp sit at TM positions
+ 0 and 2 rather than at the tail — we fall back to a greedy left-to-right
+ subsequence search through tm_sizes.
+
+ Parameters
+ ----------
+ spatial_shape : tuple of int
+ Shape of the prediction site with the sample axis removed.
+ tm_sizes : tuple of int
+ Sizes of each TensorManager dimension in order.
+
+ Returns
+ -------
+ list of int
+ tm_dims[i] is the TM dimension index for spatial_shape[i].
+ """
+ n_tm = len(tm_sizes)
+ n_sp = len(spatial_shape)
+
+ rt_dims = [n_tm - n_sp + i for i in range(n_sp)]
+
+ # Check whether non-broadcast axes agree with their right-to-left positions.
+ rt_ok = all(
+ spatial_shape[i] == tm_sizes[rt_dims[i]]
+ for i in range(n_sp)
+ if spatial_shape[i] != 1
+ )
+ if rt_ok:
+ return rt_dims
+
+ # Right-to-left is wrong: greedy left-to-right subsequence match on
+ # non-broadcast dims. Size-1 (broadcast) dims keep their rt_dims entry
+ # because modulo makes the exact TM position irrelevant.
+ result = list(rt_dims)
+ j = 0
+ for i in range(n_sp):
+ if spatial_shape[i] == 1:
+ continue
+ while j < n_tm and tm_sizes[j] != spatial_shape[i]:
+ j += 1
+ if j >= n_tm:
+ return rt_dims # No match — fall back to right-to-left
+ result[i] = j
+ j += 1
+
+ return result
+
+
+def copy_orchestrator(orchestrator,
+ t_pre=None,
+ t_sel=None,
+ titrant_conc=None,
+ genotypes=None):
"""
Generate a fresh ModelOrchestrator instance with model components from an old
ModelOrchestrator and new quantitative data passed in by the user.
- The default behavior is to use values from `model_class.growth_df`.
- Quantitative inputs (t_pre, t_sel, titrant_conc) can be expanded beyond
- those in the original dataframe.
+ The default behavior is to use values from `orchestrator.growth_df`.
+ Quantitative inputs (t_pre, t_sel, titrant_conc) can be expanded beyond
+ those in the original dataframe.
Parameters
----------
- model_class : ModelOrchestrator
+ orchestrator : ModelOrchestrator
The original ModelOrchestrator instance to copy.
t_pre : list, optional
- List of timepoints for pre-growth. Must be >= 0. If None, uses the
- value(s) from `model_class.growth_df`.
+ List of timepoints for pre-growth. Must be >= 0. If None, uses the
+ value(s) from `orchestrator.growth_df`.
t_sel : list, optional
- List of timepoints for selection. Must be >= 0. If None, uses
- values from `model_class.growth_df`.
+ List of timepoints for selection. Must be >= 0. If None, uses
+ values from `orchestrator.growth_df`.
titrant_conc : list, optional
- List of titrant concentrations. Must be >= 0. If None, uses
- values from `model_class.growth_df`.
+ List of titrant concentrations. Must be >= 0. If None, uses
+ values from `orchestrator.growth_df`.
+ genotypes : list, optional
+ Subset of genotypes to include. Must be a subset of those in
+ ``orchestrator.growth_df``. If None, uses all genotypes.
Returns
-------
ModelOrchestrator
- A new ModelOrchestrator instance initialized with the exhaustive
+ A new ModelOrchestrator instance initialized with the exhaustive
combinations of all inputs.
"""
- df = model_class.growth_df
+ df = orchestrator.growth_df
def _get_input(value, col_name):
"""
@@ -74,10 +138,23 @@ def _get_input(value, col_name):
titrant_conc_list = _get_input(titrant_conc, "titrant_conc")
# Get all unique categorical combinations
- categorical_cols = ["replicate", "condition_pre", "condition_sel",
+ categorical_cols = ["replicate", "condition_pre", "condition_sel",
"titrant_name", "genotype"]
unique_cats = df[categorical_cols].drop_duplicates()
+ # Subset genotypes when requested. Doing this here ensures the new
+ # orchestrator's TM labels differ from the original's, which causes the
+ # parameter-slicing loop in predict() to fire correctly.
+ if genotypes is not None:
+ all_genos = set(str(g) for g in unique_cats["genotype"].tolist())
+ geno_strs = [str(g) for g in genotypes]
+ missing = [g for g in geno_strs if g not in all_genos]
+ if missing:
+ raise ValueError(f"Genotype(s) not found in orchestrator: {missing}")
+ unique_cats = unique_cats[
+ unique_cats["genotype"].astype(str).isin(set(geno_strs))
+ ]
+
# Build exhaustive combinations of quantitative values
quant_combos = list(itertools.product(t_pre_list, t_sel_list, titrant_conc_list))
quant_df = pd.DataFrame(quant_combos, columns=["t_pre", "t_sel", "titrant_conc"])
@@ -91,18 +168,78 @@ def _get_input(value, col_name):
# We keep the binding_df as is, as it's keyed by genotype/titrant_name
# and we aren't subsetting those in this step.
- new_binding_df = model_class.binding_df.copy()
+ new_binding_df = orchestrator.binding_df.copy()
# Create new ModelOrchestrator using settings from the old one
- settings = model_class.settings.copy()
-
+ settings = orchestrator.settings.copy()
+
+ # When subsetting genotypes, restrict spiked_genotypes in settings to
+ # only those present in the subset. Spiked genotypes affect per-genotype
+ # masks (spiked_mask) which control per_geno_loc/scale at prediction time
+ # via the non-centred parameterisation — so any spiked genotype that IS in
+ # the subset must remain marked as spiked; those outside the subset are
+ # simply omitted. Dropping all spiked_genotypes would silently mis-classify
+ # in-subset spiked genotypes as library genotypes and corrupt predictions.
+ if genotypes is not None:
+ geno_set = set(geno_strs)
+ orig_spiked = settings.get("spiked_genotypes") or []
+ settings["spiked_genotypes"] = [
+ g for g in orig_spiked if str(g) in geno_set
+ ]
+
return ModelOrchestrator(
growth_df=new_growth_df,
binding_df=new_binding_df,
**settings
)
-def predict(model_class,
+def _convert_map_params(map_params, model_trace):
+ """
+ Convert a raw MAP parameter dict to a posterior dict suitable for Predictive.
+
+ ``AutoDelta`` guide parameters are stored in *constrained* space with
+ ``{site}_auto_loc`` keys because ``RunInference.write_params`` calls
+ ``svi.get_params(svi_state)``, which applies the SVI ``constrain_fn``
+ before saving. ``Predictive(posterior_samples=...)`` expects values keyed
+ by bare site name with a leading sample dimension. This function strips
+ the ``_auto_loc`` suffix and adds that dimension so the rest of
+ :func:`predict` can treat the MAP point as a 1-sample posterior.
+
+ .. note::
+ Do **not** apply ``biject_to`` or any other transform here. The values
+ are already constrained; a second transform would corrupt positive-
+ constrained sites (e.g. ``ln_cfu0_hyper_scale``) by applying ``exp``
+ to a value that is already in positive space.
+
+ Parameters
+ ----------
+ map_params : dict-like
+ Raw MAP parameter dict (e.g., from ``np.load("_params.npz")``).
+ Keys end with ``_auto_loc``; values are already in constrained space.
+ model_trace : dict
+ NumPyro model trace (unused; retained for interface compatibility).
+
+ Returns
+ -------
+ dict
+ Constrained values keyed by bare site name, each with shape
+ ``(1, *site_shape)``.
+ """
+ constrained = {}
+ for k, v in map_params.items():
+ k = str(k)
+ if not k.endswith("_auto_loc"):
+ continue
+ site_name = k[: -len("_auto_loc")]
+ val = jnp.array(np.asarray(v))
+ # Values are already constrained (svi.get_params() applied constrain_fn).
+ # Add the leading sample dimension expected by Predictive.
+ constrained[site_name] = jnp.expand_dims(val, 0)
+
+ return constrained
+
+
+def predict(orchestrator,
param_posteriors,
predict_sites=None,
q_to_get=None,
@@ -118,7 +255,7 @@ def predict(model_class,
Parameters
----------
- model_class : ModelOrchestrator
+ orchestrator : ModelOrchestrator
The original ModelOrchestrator used for training.
param_posteriors : dict or str
Posterior samples. Can be a dictionary, or path to .h5 file.
@@ -138,16 +275,16 @@ def predict(model_class,
quantile predictions. If None, uses all available samples.
t_pre : list, optional
List of timepoints for pre-growth. Must be >= 0. If None, uses the
- value(s) from `model_class.growth_df`.
+ value(s) from `orchestrator.growth_df`.
t_sel : list, optional
List of timepoints for selection. Must be >= 0. If None, uses
- values from `model_class.growth_df`.
+ values from `orchestrator.growth_df`.
titrant_conc : list, optional
List of titrant concentrations. Must be >= 0. If None, uses
- values from `model_class.growth_df`.
+ values from `orchestrator.growth_df`.
genotypes : list, optional
List of genotypes to include in the prediction. Must be a subset
- of those in `model_class.growth_df`. If None, uses all genotypes.
+ of those in `orchestrator.growth_df`. If None, uses all genotypes.
Returns
-------
@@ -167,44 +304,55 @@ def predict(model_class,
if isinstance(predict_sites, str):
predict_sites = [predict_sites]
- # Create the expanded prediction model
- new_mc = copy_model_class(model_class,
- t_pre=t_pre,
- t_sel=t_sel,
- titrant_conc=titrant_conc)
-
+ # Create the expanded prediction model, subsetting genotypes if requested.
+ # Passing genotypes to copy_orchestrator ensures the new orchestrator's TM
+ # labels are a strict subset of the original's, which triggers parameter
+ # slicing in the loop below. Validation of unknown genotype names is also
+ # handled inside copy_orchestrator.
+ new_orchestrator = copy_orchestrator(orchestrator,
+ t_pre=t_pre,
+ t_sel=t_sel,
+ titrant_conc=titrant_conc,
+ genotypes=genotypes)
+
+ # After copy_orchestrator the new TM contains exactly the requested
+ # genotypes (or all genotypes when genotypes=None).
+ genotypes = new_orchestrator.growth_tm.tensor_dim_labels[-1].tolist()
# -------------------------------------------------------------------------
- # Genotype subsetting
-
- all_genotypes = new_mc.growth_tm.tensor_dim_labels[-1]
- if genotypes is None:
- genotypes = all_genotypes.tolist()
- genotype_indices = np.arange(len(all_genotypes))
- else:
- # Validate and find indices
- genotype_indices = []
- for g in genotypes:
- matches = np.where(all_genotypes == g)[0]
- if len(matches) == 0:
- # Try string fallback
- matches = np.where(all_genotypes.astype(str) == str(g))[0]
-
- if len(matches) == 0:
- raise ValueError(f"Genotype {g} not found in model.")
- genotype_indices.append(matches[0])
- genotype_indices = np.array(genotype_indices)
+ # Model trace — run first so bijections are available for MAP detection.
+ #
+ # We use the original orchestrator because posteriors match its structure.
+
+ seeded_model = seed(orchestrator.jax_model, rng_seed=0)
+ traced_model = trace(seeded_model)
+ model_trace = traced_model.get_trace(data=orchestrator.data,
+ priors=orchestrator.priors)
- # Use existing get_batch to subset the data tensors
- subset_data = new_mc.get_batch(new_mc.data, jnp.array(genotype_indices))
+ # -------------------------------------------------------------------------
+ # MAP param conversion (if needed)
+ #
+ # Raw MAP params from RunInference.write_params() / np.load("_params.npz")
+ # have keys like ``{site}_auto_loc`` in *constrained* space (because
+ # write_params receives the output of svi.get_params(), which applies
+ # constrain_fn) and no leading sample dimension. Rename the keys and add
+ # a size-1 leading dim so the rest of this function treats them as a
+ # 1-sample posterior.
+ #
+ # Genuine posterior files produced by get_posteriors() / get_map_posteriors()
+ # already store constrained values keyed by bare site name and are left
+ # unchanged.
+
+ if any(str(k).endswith("_auto_loc") for k in param_posteriors.keys()):
+ param_posteriors = _convert_map_params(param_posteriors, model_trace)
# -------------------------------------------------------------------------
# Sample selection from posterior
-
+
# Identify how many samples we have
first_key = next(iter(param_posteriors.keys()))
total_available = param_posteriors[first_key].shape[0]
-
+
# Sample indices for running through the model (quantile computation)
n_for_quantiles = total_available if num_marginal_samples is None else min(num_marginal_samples, total_available)
rng = np.random.default_rng()
@@ -213,13 +361,6 @@ def predict(model_class,
# -------------------------------------------------------------------------
# Parameter slicing
-
- # Run a trace of the original model to identify plate structure
- # We use the original model class because posteriors match its structure.
- seeded_model = seed(model_class.jax_model, rng_seed=0)
- traced_model = trace(seeded_model)
- model_trace = traced_model.get_trace(data=model_class.data,
- priors=model_class.priors)
sliced_samples = {}
for site_name, site in model_trace.items():
@@ -244,14 +385,14 @@ def predict(model_class,
# Find the corresponding dimension in the TensorManager
dim_idx = None
- for i, name in enumerate(model_class.growth_tm.tensor_dim_names):
+ for i, name in enumerate(orchestrator.growth_tm.tensor_dim_names):
if name.lower() in plate_name:
dim_idx = i
break
if dim_idx is not None:
- old_labels = model_class.growth_tm.tensor_dim_labels[dim_idx]
- new_labels = new_mc.growth_tm.tensor_dim_labels[dim_idx]
+ old_labels = orchestrator.growth_tm.tensor_dim_labels[dim_idx]
+ new_labels = new_orchestrator.growth_tm.tensor_dim_labels[dim_idx]
if not np.array_equal(old_labels, new_labels):
try:
@@ -270,64 +411,74 @@ def predict(model_class,
# -------------------------------------------------------------------------
# Run Prediction
- predictive = Predictive(new_mc.jax_model,
- posterior_samples=sliced_samples,
+ predictive = Predictive(new_orchestrator.jax_model,
+ posterior_samples=sliced_samples,
return_sites=predict_sites)
-
+
# We need a key, even if not used for randomness in deterministic sites
- predict_key = jax.random.PRNGKey(0)
- predictions = predictive(predict_key,
- data=subset_data,
- priors=new_mc.priors)
-
+ predict_key = jax.random.PRNGKey(0)
+
+ # Build a canonical-indexed data view for prediction. The static
+ # new_orchestrator.data.growth.batch_idx uses a binding-first ordering
+ # (binding genotypes at positions 0..num_binding-1, then library genotypes
+ # at positions num_binding..N-1). Posterior samples, however, are always
+ # stored in canonical genotype order (index g = genotype g). Passing the
+ # static data to Predictive would misalign posterior params vs. mask
+ # lookups for all non-canonical positions.
+ #
+ # get_batch with arange(num_genotype) sets batch_idx = arange(num_genotype)
+ # (canonical), which is exactly what get_posteriors uses when saving
+ # samples. This makes per-genotype lookups (masks, offsets) consistent
+ # with the stored posterior values.
+ num_geno = new_orchestrator.data.growth.num_genotype
+ all_indices = jnp.arange(num_geno, dtype=jnp.int32)
+ pred_data = _get_batch(new_orchestrator.data, all_indices)
+
+ # Use the ORIGINAL orchestrator's priors so that any pinned hyperpriors
+ # (e.g. activity_hyper_loc/activity_hyper_scale in the calibration model)
+ # remain pinned at their calibrated values rather than being sampled from
+ # the default prior. new_orchestrator is rebuilt from settings-only
+ # (component names, no priors) so its priors are all-default.
+ predictions = predictive(predict_key,
+ data=pred_data,
+ priors=orchestrator.priors)
+
# -------------------------------------------------------------------------
# Calculate Quantiles and Join
-
- # tm._pivot_index columns in df contain the integer codes for each dimension
- # (replicate_idx, time_idx, etc.)
- # We want a dataframe that only has the subsetted genotypes.
- base_df = new_mc.growth_df.copy()
- base_df = base_df[base_df["genotype"].isin(genotypes)].copy()
+
+ # new_orchestrator already contains exactly the requested genotypes, so
+ # growth_df needs no further filtering.
+ base_df = new_orchestrator.growth_df.copy()
# Replace the dummy ln_cfu/ln_cfu_std zeros with observed values from the
- # original model_class.growth_df (NaN where there is no matching observation,
+ # original orchestrator.growth_df (NaN where there is no matching observation,
# e.g. for expanded prediction grids).
merge_keys = ["replicate", "condition_pre", "condition_sel",
"titrant_name", "genotype", "t_pre", "t_sel", "titrant_conc"]
obs_cols = merge_keys + ["ln_cfu", "ln_cfu_std"]
- orig_obs = model_class.growth_df[obs_cols].drop_duplicates(subset=merge_keys)
+ orig_obs = orchestrator.growth_df[obs_cols].drop_duplicates(subset=merge_keys)
base_df = base_df.drop(columns=["ln_cfu", "ln_cfu_std"]).merge(
orig_obs, on=merge_keys, how="left"
)
-
- # Re-calculate indices for the subsetted dataframe relative to the
- # using the new_mc TM but subset the df.
- tm = new_mc.growth_tm
+
+ tm = new_orchestrator.growth_tm
indices = [base_df[f"{dim}_idx"].values for dim in tm.tensor_dim_names]
-
- # Since we used get_batch, the predictions only have the subsetted genotypes.
- # The last dimension of predictions[site] will match the length of genotype_indices.
- # To map back to the tensor indices, we need to map genotype_idx in base_df
- # to its relative position in the genotype_indices array.
-
- # Mapping from original genotype index to new relative index
- geno_map = {orig_idx: i for i, orig_idx in enumerate(genotype_indices)}
- relative_geno_indices = base_df["genotype_idx"].map(geno_map).values
-
- # Update indices for genotype to be relative for the prediction tensor
- indices[-1] = relative_geno_indices
+
+ # TM dimension sizes matching the current (possibly genotype-subsetted) indices.
+ tm_sizes = tuple(int(idx.max()) + 1 if len(idx) > 0 else 1 for idx in indices)
all_dfs = {}
for site in predict_sites:
site_samples = predictions[site] # shape: (num_samples, ...)
- # Compute the per-row index tuple once — shape is the same for every
- # quantile and for individual sample draws.
- # q_arr drops the leading sample axis, so spatial_shape == site_samples.shape[1:].
- # We align right-to-left and use modulo to handle size-1 broadcast axes.
+ # Map each axis of the site's spatial shape to the correct TM dimension.
+ # Right-to-left alignment works for full-rank and tail-aligned sites;
+ # _align_site_to_tm_dims falls back to subsequence matching for sites
+ # like ln_cfu0 whose batch dims are not at the tail of the TM ordering.
spatial_shape = site_samples.shape[1:]
+ tm_dims = _align_site_to_tm_dims(spatial_shape, tm_sizes)
aligned_indices = tuple(
- indices[len(indices) - len(spatial_shape) + i] % spatial_shape[i]
+ indices[tm_dims[i]] % spatial_shape[i]
for i in range(len(spatial_shape))
)
@@ -341,7 +492,7 @@ def predict(model_class,
has_broadcast = any(sz == 1 for sz in spatial_shape)
if has_broadcast:
nontrivial_tm_idx = [
- len(indices) - len(spatial_shape) + i
+ tm_dims[i]
for i, sz in enumerate(spatial_shape) if sz > 1
]
if nontrivial_tm_idx:
diff --git a/src/tfscreen/tfmodel/analysis/prior_predictive.py b/src/tfscreen/tfmodel/analysis/prior_predictive.py
index 65521cb6..2e46c66f 100644
--- a/src/tfscreen/tfmodel/analysis/prior_predictive.py
+++ b/src/tfscreen/tfmodel/analysis/prior_predictive.py
@@ -6,7 +6,7 @@
from tfscreen.tfmodel.analysis.prediction import predict
-def draw_prior(model_class, rng_key=0, num_draws=1):
+def draw_prior(orchestrator, rng_key=0, num_draws=1):
"""
Draw samples from the model prior (no data conditioning).
@@ -17,7 +17,7 @@ def draw_prior(model_class, rng_key=0, num_draws=1):
Parameters
----------
- model_class : ModelOrchestrator
+ orchestrator : ModelOrchestrator
Fully initialised model, as returned by ``read_configuration``.
rng_key : int or jax.random.PRNGKey, optional
Random seed. An ``int`` is converted to a PRNGKey. Default 0.
@@ -41,8 +41,8 @@ def draw_prior(model_class, rng_key=0, num_draws=1):
rng_key = jax.random.PRNGKey(rng_key)
# Trace the model once to classify all sites.
- seeded = seed(model_class.jax_model, rng_seed=0)
- tr = trace(seeded).get_trace(data=model_class.data, priors=model_class.priors)
+ seeded = seed(orchestrator.jax_model, rng_seed=0)
+ tr = trace(seeded).get_trace(data=orchestrator.data, priors=orchestrator.priors)
observed = {
name for name, site in tr.items()
@@ -56,12 +56,12 @@ def draw_prior(model_class, rng_key=0, num_draws=1):
return_sites = sorted(set(tr.keys()) - observed)
predictive = Predictive(
- model_class.jax_model,
+ orchestrator.jax_model,
posterior_samples=None,
num_samples=num_draws,
return_sites=return_sites,
)
- raw = predictive(rng_key, data=model_class.data, priors=model_class.priors)
+ raw = predictive(rng_key, data=orchestrator.data, priors=orchestrator.priors)
predictions = {k: np.array(v) for k, v in raw.items() if k in deterministic}
latent_params = {k: np.array(v) for k, v in raw.items() if k not in deterministic}
@@ -69,23 +69,23 @@ def draw_prior(model_class, rng_key=0, num_draws=1):
return predictions, latent_params
-def growth_df_from_prior(model_class, latent_params, draw_idx=0, noise_rng=None):
+def growth_df_from_prior(orchestrator, latent_params, draw_idx=0, noise_rng=None):
"""
Build a synthetic growth DataFrame from one prior draw.
Maps ``latent_params`` through the forward model to obtain the
deterministic ``growth_pred`` values, then writes them into a copy of
- ``model_class.growth_df``. The result has the same rows and structure
+ ``orchestrator.growth_df``. The result has the same rows and structure
as the original DataFrame and is ready for ``tfs-fit-model``.
- Noise is added using ``model_class.growth_df['ln_cfu_std']`` as the
+ Noise is added using ``orchestrator.growth_df['ln_cfu_std']`` as the
per-observation standard deviation. This preserves the original
measurement-error structure while replacing the signal with synthetic
predictions.
Parameters
----------
- model_class : ModelOrchestrator
+ orchestrator : ModelOrchestrator
latent_params : dict
Output of ``draw_prior``. Values have shape ``(num_draws, ...)``.
draw_idx : int, optional
@@ -97,7 +97,7 @@ def growth_df_from_prior(model_class, latent_params, draw_idx=0, noise_rng=None)
Returns
-------
pandas.DataFrame
- Copy of ``model_class.growth_df`` with ``ln_cfu`` replaced by
+ Copy of ``orchestrator.growth_df`` with ``ln_cfu`` replaced by
synthetic values.
"""
# Slice to single draw so predict() treats it as one "posterior sample".
@@ -109,20 +109,20 @@ def growth_df_from_prior(model_class, latent_params, draw_idx=0, noise_rng=None)
]
pred_df = predict(
- model_class,
+ orchestrator,
single_draw,
predict_sites=["growth_pred"],
- q_to_get={"_prior_ln_cfu": 0.5},
+ q_to_get=[0.5],
num_samples=None,
)
# pred_df may contain extra rows (expanded grid); merge to keep original rows.
- out_df = model_class.growth_df.drop(columns=["ln_cfu"]).merge(
- pred_df[merge_keys + ["_prior_ln_cfu"]],
+ out_df = orchestrator.growth_df.drop(columns=["ln_cfu"]).merge(
+ pred_df[merge_keys + ["q0.5"]],
on=merge_keys,
how="left",
)
- out_df = out_df.rename(columns={"_prior_ln_cfu": "ln_cfu"})
+ out_df = out_df.rename(columns={"q0.5": "ln_cfu"})
if noise_rng is not None:
out_df["ln_cfu"] = (
diff --git a/src/tfscreen/tfmodel/analysis/sbc.py b/src/tfscreen/tfmodel/analysis/sbc.py
deleted file mode 100644
index 1a64bb2e..00000000
--- a/src/tfscreen/tfmodel/analysis/sbc.py
+++ /dev/null
@@ -1,331 +0,0 @@
-"""
-Simulation-Based Calibration (SBC) rank statistics.
-
-Theory
-------
-For a correctly specified model, the posterior rank of the true parameter
-value should be uniformly distributed on [0, 1]. Given N ground-truth /
-posterior pairs (each drawn from the same generative model), we:
-
- 1. For each parameter array θ of shape (*D), compute the rank of the
- ground-truth value θ_true among the S posterior samples::
-
- rank[...] = mean(posterior_samples < θ_true, axis=0) ∈ [0, 1]
-
- 2. Pool all ranks across runs and parameter elements.
- 3. Test uniformity with a KS test and summarise with a rank histogram.
-
-References
-----------
-Talts et al. (2018) "Validating Bayesian Inference Algorithms with
-Simulation-Based Calibration". arXiv:1804.06788.
-"""
-
-import glob
-import os
-import warnings
-from typing import Dict, List, Optional, Tuple
-
-import h5py
-import numpy as np
-import pandas as pd
-from matplotlib import pyplot as plt
-from scipy.stats import ks_1samp, uniform
-
-
-# ---------------------------------------------------------------------------
-# Low-level helpers
-# ---------------------------------------------------------------------------
-
-def _load_h5_params(path: str) -> Dict[str, np.ndarray]:
- """Load all dataset arrays from an HDF5 file."""
- out = {}
- with h5py.File(path, "r") as hf:
- for key in hf.keys():
- out[key] = hf[key][:]
- return out
-
-
-def compute_sbc_ranks(ground_truth_path: str,
- posterior_path: str) -> Dict[str, np.ndarray]:
- """
- Compute the posterior rank of each ground-truth parameter value.
-
- For each parameter present in *both* files the rank is the fraction of
- posterior samples that are strictly less than the ground-truth value::
-
- rank[i] = mean(posterior_samples[:, i] < gt_value[i]) ∈ [0, 1]
-
- Parameters
- ----------
- ground_truth_path : str
- HDF5 file produced by ``tfs-sample-prior``. Each dataset has shape
- ``(1, *param_shape)`` (num_draws=1).
- posterior_path : str
- HDF5 file produced by ``tfs-sample-posterior``. Each dataset has
- shape ``(S, *param_shape)`` with S posterior samples.
-
- Returns
- -------
- dict
- Mapping from parameter name to a 1-D rank array (one value per
- scalar element of the parameter). Only parameters present in both
- files are included.
- """
- gt_params = _load_h5_params(ground_truth_path)
- post_params = _load_h5_params(posterior_path)
-
- ranks = {}
- for name, gt_val in gt_params.items():
- if name not in post_params:
- continue
-
- # Ground truth: shape (1, *param_shape) → squeeze to (*param_shape)
- gt = gt_val[0] if gt_val.ndim >= 1 else gt_val
- post = post_params[name] # (S, *param_shape)
-
- if post.ndim < 1 or post.shape[0] == 0:
- continue
-
- try:
- gt_f = np.asarray(gt, dtype=float)
- post_f = np.asarray(post, dtype=float)
- except (TypeError, ValueError):
- continue # skip non-numeric parameters
-
- # Broadcast: (S, *param_shape) vs (*param_shape)
- rank = np.mean(post_f < gt_f, axis=0).ravel() # shape: (n_elements,)
- ranks[name] = rank
-
- return ranks
-
-
-# ---------------------------------------------------------------------------
-# Directory scanner
-# ---------------------------------------------------------------------------
-
-def _find_pairs(sbc_dir: str) -> List[Tuple[str, str, Optional[str]]]:
- """
- Find (run_id, ground_truth_path, posterior_path_or_None) triples.
-
- Scans *sbc_dir* for files matching ``*_ground_truth.h5``. For each,
- looks for a corresponding ``*_posterior.h5`` with the same prefix
- (replacing ``_ground_truth`` with ``_posterior``).
- """
- gt_files = sorted(glob.glob(os.path.join(sbc_dir, "*_ground_truth.h5")))
- pairs = []
- for gt_path in gt_files:
- basename = os.path.basename(gt_path)
- run_id = basename[: -len("_ground_truth.h5")]
- post_path = os.path.join(sbc_dir, f"{run_id}_posterior.h5")
- pairs.append((run_id, gt_path,
- post_path if os.path.exists(post_path) else None))
- return pairs
-
-
-# ---------------------------------------------------------------------------
-# Main summary function
-# ---------------------------------------------------------------------------
-
-def summarize_sbc(sbc_dir: str, out_prefix: Optional[str] = None) -> pd.DataFrame:
- """
- Scan *sbc_dir* for SBC run pairs and compute rank-uniformity statistics.
-
- For each ``*_ground_truth.h5`` / ``*_posterior.h5`` pair found, compute
- the posterior rank of every ground-truth parameter value. Pool ranks
- across runs and write three output files:
-
- ``{out_prefix}_sbc_summary.csv``
- One row per parameter with mean rank, KS statistic, and p-value.
- Mean rank near 0.5 and high p-value indicate good calibration.
-
- ``{out_prefix}_sbc_ranks.csv``
- Long-form table of every individual rank value, for custom analysis.
- Columns: ``run_id``, ``param``, ``element_idx``, ``rank``.
-
- ``{out_prefix}_rank_hist.pdf``
- One-panel-per-parameter rank histogram. A flat histogram indicates
- uniform ranks (good calibration).
-
- Parameters
- ----------
- sbc_dir : str
- Directory containing ``*_ground_truth.h5`` and ``*_posterior.h5``
- files produced by ``tfs-sample-prior`` and ``tfs-sample-posterior``.
- out_prefix : str, optional
- Prefix for all output files. Defaults to ``{sbc_dir}/sbc``.
-
- Returns
- -------
- pd.DataFrame
- Per-parameter summary table (same content as the CSV).
- """
- sbc_dir = os.path.abspath(sbc_dir)
- if not os.path.isdir(sbc_dir):
- raise FileNotFoundError(f"SBC directory not found: {sbc_dir}")
-
- if out_prefix is None:
- out_prefix = os.path.join(sbc_dir, "sbc")
-
- pairs = _find_pairs(sbc_dir)
- if not pairs:
- print(f"No *_ground_truth.h5 files found in {sbc_dir}/", flush=True)
- return pd.DataFrame()
-
- n_total = len(pairs)
- n_posterior = sum(1 for _, _, p in pairs if p is not None)
- print(
- f"Found {n_total} ground-truth file(s), "
- f"{n_posterior} with matching posterior.",
- flush=True,
- )
- if n_posterior == 0:
- warnings.warn(
- "No posterior files found. Run tfs-sample-posterior for each "
- "synthetic dataset before calling tfs-summarize-sbc."
- )
- return pd.DataFrame()
-
- # Accumulate ranks: param_name → list of rank arrays (one per run)
- all_ranks: Dict[str, List[np.ndarray]] = {}
- rank_rows = []
-
- for run_id, gt_path, post_path in pairs:
- if post_path is None:
- continue
- try:
- run_ranks = compute_sbc_ranks(gt_path, post_path)
- except Exception as exc:
- warnings.warn(f"Skipping run '{run_id}': {exc}")
- continue
-
- for param, rank_arr in run_ranks.items():
- all_ranks.setdefault(param, []).append(rank_arr)
- for idx, r in enumerate(rank_arr):
- rank_rows.append({
- "run_id": run_id,
- "param": param,
- "element_idx": idx,
- "rank": float(r),
- })
-
- if not all_ranks:
- warnings.warn("No matching parameters found across ground-truth and posterior files.")
- return pd.DataFrame()
-
- # -----------------------------------------------------------------
- # Per-parameter summary: mean rank, KS test against Uniform(0,1)
- # -----------------------------------------------------------------
- summary_rows = []
- for param in sorted(all_ranks):
- pooled = np.concatenate(all_ranks[param])
- pooled = pooled[np.isfinite(pooled)]
- n_runs = len(all_ranks[param])
- n_vals = len(pooled)
- if n_vals == 0:
- continue
- mean_rank = float(np.mean(pooled))
- ks_stat, ks_pval = ks_1samp(pooled, uniform(0, 1).cdf)
- summary_rows.append({
- "param": param,
- "n_runs": n_runs,
- "n_ranks": n_vals,
- "mean_rank": mean_rank,
- "expected_rank": 0.5,
- "ks_stat": float(ks_stat),
- "ks_pval": float(ks_pval),
- })
-
- summary_df = pd.DataFrame(summary_rows)
-
- # -----------------------------------------------------------------
- # Write outputs
- # -----------------------------------------------------------------
- summary_csv = f"{out_prefix}_sbc_summary.csv"
- summary_df.to_csv(summary_csv, index=False)
- print(f"Wrote per-parameter summary → {summary_csv}", flush=True)
-
- if rank_rows:
- ranks_csv = f"{out_prefix}_sbc_ranks.csv"
- ranks_df = pd.DataFrame(rank_rows)
- ranks_df.to_csv(ranks_csv, index=False)
- print(f"Wrote all rank values → {ranks_csv}", flush=True)
- else:
- ranks_df = pd.DataFrame()
-
- _plot_rank_histograms(all_ranks, out_prefix)
-
- print(
- f"\nSummary: {len(summary_df)} parameter(s) across "
- f"{n_posterior} run(s).",
- flush=True,
- )
- low_pval = summary_df[summary_df["ks_pval"] < 0.05] if len(summary_df) else pd.DataFrame()
- if len(low_pval):
- print(
- f" {len(low_pval)} parameter(s) with KS p-value < 0.05 "
- f"(possible miscalibration):",
- flush=True,
- )
- for _, row in low_pval.iterrows():
- print(
- f" {row['param']:40s} KS={row['ks_stat']:.3f} p={row['ks_pval']:.3g}",
- flush=True,
- )
- else:
- print(" All parameters pass KS uniformity test (p ≥ 0.05).", flush=True)
-
- return summary_df
-
-
-# ---------------------------------------------------------------------------
-# Plotting
-# ---------------------------------------------------------------------------
-
-_HIST_BINS = 10
-
-
-def _plot_rank_histograms(all_ranks: Dict[str, List[np.ndarray]],
- out_prefix: str) -> None:
- """Write a PDF with one rank-histogram panel per parameter."""
- params = sorted(all_ranks)
- n = len(params)
- if n == 0:
- return
-
- ncols = min(4, n)
- nrows = (n + ncols - 1) // ncols
- fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3 * nrows),
- squeeze=False)
-
- for ax_idx, param in enumerate(params):
- ax = axes[ax_idx // ncols][ax_idx % ncols]
- pooled = np.concatenate(all_ranks[param])
- pooled = pooled[np.isfinite(pooled)]
- n_vals = len(pooled)
-
- ax.hist(pooled, bins=_HIST_BINS, range=(0, 1),
- density=True, color="steelblue", alpha=0.75, edgecolor="white")
- ax.axhline(1.0, color="firebrick", lw=1.5, ls="--", label="uniform")
- ax.set_xlim(0, 1)
- ax.set_xlabel("Posterior rank")
- ax.set_ylabel("Density")
-
- # Truncate long param names for the title
- title = param if len(param) <= 30 else f"…{param[-27:]}"
- _, ks_pval = ks_1samp(pooled, uniform(0, 1).cdf)
- ax.set_title(f"{title}\n(n={n_vals}, p={ks_pval:.2g})", fontsize=9)
-
- # Hide unused panels
- for ax_idx in range(n, nrows * ncols):
- axes[ax_idx // ncols][ax_idx % ncols].set_visible(False)
-
- fig.tight_layout()
- pdf_path = f"{out_prefix}_rank_hist.pdf"
- try:
- fig.savefig(pdf_path)
- print(f"Wrote rank histograms → {pdf_path}", flush=True)
- except Exception as exc:
- warnings.warn(f"Could not write rank histogram PDF: {exc}")
- finally:
- plt.close(fig)
diff --git a/src/tfscreen/tfmodel/configuration_io.py b/src/tfscreen/tfmodel/configuration_io.py
index a364195e..06f6752d 100644
--- a/src/tfscreen/tfmodel/configuration_io.py
+++ b/src/tfscreen/tfmodel/configuration_io.py
@@ -5,6 +5,16 @@
import jax.numpy as jnp
import warnings
from tfscreen.__version__ import __version__
+from tfscreen.util.validation import check_unknown_keys
+
+# All recognized top-level keys for a tfmodel config file.
+TFMODEL_KNOWN_KEYS = frozenset({
+ "tfscreen_version",
+ "data",
+ "components",
+ "priors_file",
+ "guesses_file",
+})
def _extract_scalars(obj, prefix=""):
"""Recursively extract scalar values from PriorsClass."""
@@ -103,16 +113,17 @@ def _update_dataclass(dc, prefix, flat_dict):
return dataclasses.replace(dc, **updates)
return dc
-def write_configuration(gm,
+def write_configuration(orchestrator,
out_prefix,
growth_df_path=None,
- binding_df_path=None):
+ binding_df_path=None,
+ presplit_df_path=None):
"""
Write model configuration and extracted priors/guesses to files.
Parameters
----------
- gm : ModelOrchestrator
+ orchestrator : ModelOrchestrator
Initialized ModelOrchestrator object.
out_prefix : str
Root filename for output files.
@@ -126,12 +137,12 @@ def write_configuration(gm,
guesses_list = []
# Extract priors
- priors_raw = _extract_scalars(gm.priors)
+ priors_raw = _extract_scalars(orchestrator.priors)
for k, v in priors_raw.items():
priors_list.append(pd.DataFrame({"parameter": [k], "value": [v]}))
# Extract guesses
- for k, v in gm.init_params.items():
+ for k, v in orchestrator.init_params.items():
if not hasattr(v, 'shape') or len(v.shape) == 0:
guesses_list.append(pd.DataFrame({"parameter": [k], "value": [v]}))
@@ -144,11 +155,13 @@ def write_configuration(gm,
data_paths["growth"] = growth_df_path
if binding_df_path is not None:
data_paths["binding"] = binding_df_path
+ if presplit_df_path is not None:
+ data_paths["presplit"] = presplit_df_path
config = {
"tfscreen_version": __version__,
"data": data_paths,
- "components": gm.settings,
+ "components": orchestrator.settings,
"priors_file": os.path.basename(priors_path),
"guesses_file": os.path.basename(guesses_path)
}
@@ -159,7 +172,7 @@ def write_configuration(gm,
print(f"Wrote configuration to {yaml_path}")
# Process array guesses and any others
- tm = gm.growth_tm
+ tm = orchestrator.growth_tm
if tm is not None:
cond_rep_map = tm.map_groups.get("condition_rep", pd.DataFrame())
geno_map = tm.map_groups.get("genotype", pd.DataFrame())
@@ -171,7 +184,7 @@ def write_configuration(gm,
theta_map = pd.DataFrame()
ln_cfu0_map = pd.DataFrame()
- for k, v in gm.init_params.items():
+ for k, v in orchestrator.init_params.items():
if not hasattr(v, 'shape') or len(v.shape) == 0:
continue
@@ -248,7 +261,7 @@ def read_configuration(config_file):
Returns
-------
- gm : ModelOrchestrator
+ orchestrator : ModelOrchestrator
Initialized ModelOrchestrator object.
init_params : dict
Dictionary of initial parameters for the model.
@@ -260,10 +273,13 @@ def read_configuration(config_file):
with open(config_file, "r") as f:
config = yaml.safe_load(f)
+ check_unknown_keys(config, TFMODEL_KNOWN_KEYS, label="tfmodel config")
+
# Check sanity of format and read in data paths and components
if "data" in config and "components" in config:
growth_df_path = config["data"].get("growth")
binding_df_path = config["data"].get("binding")
+ presplit_df_path = config["data"].get("presplit")
settings = config["components"]
else:
raise ValueError(f"Configuration file '{config_file}' has an unrecognized format.")
@@ -273,10 +289,14 @@ def read_configuration(config_file):
warnings.warn(f"Configuration file version {config['tfscreen_version']} does not match current tfscreen version {__version__}")
batch_size = settings.pop("batch_size", None)
+ # presplit_df is a data path, not a component setting; pop it if it
+ # ended up in settings (older configs that serialised it there).
+ settings.pop("presplit_df", None)
- gm = ModelOrchestrator(growth_df_path,
+ orchestrator = ModelOrchestrator(growth_df_path,
binding_df_path,
batch_size=batch_size,
+ presplit_df=presplit_df_path,
**settings)
# Update Priors from CSV
@@ -296,8 +316,8 @@ def read_configuration(config_file):
except (ValueError, TypeError):
flat_priors[k] = v
- new_priors = _update_dataclass(gm.priors, "", flat_priors)
- gm._priors = new_priors
+ new_priors = _update_dataclass(orchestrator.priors, "", flat_priors)
+ orchestrator._priors = new_priors
# Construct init_params from CSV
guesses_file = config.get("guesses_file")
@@ -316,10 +336,19 @@ def read_configuration(config_file):
sorted_group = df_group.sort_values("flat_index") if "flat_index" in df_group else df_group
val_array = sorted_group["value"].values
- if param_name in gm.init_params:
- orig_val = gm.init_params[param_name]
+ if param_name in orchestrator.init_params:
+ orig_val = orchestrator.init_params[param_name]
if hasattr(orig_val, 'shape') and orig_val.shape != ():
orig_shape = orig_val.shape
+ if val_array.size != np.prod(orig_shape):
+ raise ValueError(
+ f"Parameter '{param_name}' has {val_array.size} "
+ f"{'value' if val_array.size == 1 else 'values'} in "
+ f"'{guesses_file}' but the current model expects shape "
+ f"{orig_shape} ({int(np.prod(orig_shape))} values). "
+ f"The guesses file is likely stale — regenerate it with "
+ f"tfs-configure-model and re-run tfs-prefit-calibration."
+ )
init_params[param_name] = jnp.array(val_array.reshape(orig_shape))
else:
init_params[param_name] = float(val_array[0])
@@ -327,11 +356,11 @@ def read_configuration(config_file):
init_params[param_name] = float(df_group["value"].iloc[0])
missing_params = []
- for k in gm.init_params.keys():
+ for k in orchestrator.init_params.keys():
if k not in init_params:
missing_params.append(k)
if len(missing_params) > 0:
raise ValueError(f"Missing initial guesses for parameters: {missing_params}")
- return gm, init_params
+ return orchestrator, init_params
diff --git a/src/tfscreen/tfmodel/data_class.py b/src/tfscreen/tfmodel/data_class.py
index 345a67dc..cc0e963f 100644
--- a/src/tfscreen/tfmodel/data_class.py
+++ b/src/tfscreen/tfmodel/data_class.py
@@ -149,6 +149,33 @@ class BindingData:
struct_contact_distances: Any = field(pytree_node=False, default=None)
+@dataclass(frozen=True)
+class PreSplitData:
+ """
+ Holds the optional pre-split (t = -t_pre) sequencing observations.
+
+ These are taken from a single pooled aliquot just before the culture is
+ split into titrant-concentration conditions, so they carry no
+ condition_sel or titrant_conc index. The prediction for each observation
+ is simply ``ln_cfu0[replicate, condition_pre, genotype]``, making this a
+ direct side-channel constraint on the initial-population parameter.
+
+ Tensors are shaped ``(num_replicate, num_condition_pre, num_genotype)``
+ and use the same categorical orderings as the companion GrowthData so
+ that slicing with ``data.growth.batch_idx`` aligns the genotype axis.
+ """
+
+ # Data tensors
+ ln_cfu_t0: jnp.ndarray
+ ln_cfu_t0_std: jnp.ndarray
+ good_mask: jnp.ndarray
+
+ # Tensor dimensions (aligned with GrowthData)
+ num_replicate: int = field(pytree_node=False)
+ num_condition_pre: int = field(pytree_node=False)
+ num_genotype: int = field(pytree_node=False)
+
+
@dataclass(frozen=True)
class DataClass:
"""
@@ -157,17 +184,19 @@ class DataClass:
"""
num_genotype: int = field(pytree_node=False)
-
+
batch_idx: jnp.ndarray
batch_size: int = field(pytree_node=False)
not_binding_idx: jnp.ndarray
not_binding_batch_size: int = field(pytree_node=False)
- num_binding: int = field(pytree_node=False)
+ num_binding: int = field(pytree_node=False)
# GrowthData when running the joint growth+binding model; None in binding_only mode.
growth: Any = field(default=None)
binding: BindingData = field(default=None)
+ # Optional pre-split observations (t = -t_pre aliquot before condition split).
+ presplit: Any = field(default=None)
@dataclass(frozen=True)
diff --git a/src/tfscreen/tfmodel/generative/components/growth/linear.py b/src/tfscreen/tfmodel/generative/components/growth/linear.py
index 2f687ff5..73c10503 100644
--- a/src/tfscreen/tfmodel/generative/components/growth/linear.py
+++ b/src/tfscreen/tfmodel/generative/components/growth/linear.py
@@ -2,8 +2,8 @@
import jax.numpy as jnp
import numpyro as pyro
import numpyro.distributions as dist
-from flax.struct import dataclass
-from typing import Tuple
+from flax.struct import dataclass, field
+from typing import List, Optional, Tuple
from tfscreen.tfmodel.data_class import (
GrowthData
@@ -39,12 +39,83 @@ class ModelPriors:
m_loc : float
Mean of the Normal prior on the per-condition occupancy slope m.
m_scale : float
- Standard deviation of the Normal prior on m.
+ Standard deviation of the Normal prior on m (used for all conditions
+ when ``m_is_selection`` is ``None``).
+ m_scale_minus : float
+ Prior scale on m for control ('-') conditions, where theta is not
+ expected to meaningfully affect growth. Applied per-condition when
+ ``m_is_selection`` is provided.
+ m_scale_plus : float
+ Prior scale on m for selection ('+') conditions, where theta drives
+ differential growth. Applied per-condition when ``m_is_selection``
+ is provided.
+ m_is_selection : tuple of bool or None
+ Length-``num_condition_rep`` tuple; ``True`` for selection ('+')
+ conditions, ``False`` for control ('-') conditions. ``None``
+ disables the per-condition logic and uses ``m_scale`` for all
+ conditions (backward-compatible default).
"""
k_loc: float
k_scale: float
m_loc: float
m_scale: float
+ m_scale_minus: float
+ m_scale_plus: float
+ m_is_selection: tuple = field(pytree_node=False, default=None)
+
+
+# ---------------------------------------------------------------------------
+# Condition-label parsing
+# ---------------------------------------------------------------------------
+
+def _parse_condition_label(label: str) -> bool:
+ """
+ Classify one condition as selection ('+') or control ('-').
+
+ Parameters
+ ----------
+ label : str
+ Condition name as it appears in the growth DataFrame.
+
+ Returns
+ -------
+ bool
+ ``True`` if this is a selection ('+') condition,
+ ``False`` if it is a control ('-') condition.
+
+ Raises
+ ------
+ ValueError
+ If the label contains neither '+' nor '-', or contains both.
+ The error includes a detailed explanation of the naming convention.
+ """
+ has_plus = '+' in label
+ has_minus = '-' in label
+
+ if has_plus and not has_minus:
+ return True
+ if has_minus and not has_plus:
+ return False
+
+ # Both or neither: ambiguous.
+ raise ValueError(
+ f"Condition name '{label}' does not unambiguously identify whether "
+ "it is a selection or control condition.\n\n"
+ "The linear growth component uses the presence of '+' or '-' in the "
+ "condition name to apply different priors on m (the theta-to-growth "
+ "slope):\n"
+ " '+' conditions — selection media (e.g. kanR+kan, pheS+4CP): theta "
+ "meaningfully suppresses or promotes growth, so m gets the full "
+ "Normal(0, m_scale_plus) prior.\n"
+ " '-' conditions — control/no-selection media (e.g. kanR-kan, "
+ "pheS-4CP): theta has negligible effect on growth, so m gets the tight "
+ "Normal(0, m_scale_minus) prior, preventing the optimizer from "
+ "incorrectly absorbing selection-phase theta signal into the control "
+ "condition.\n\n"
+ "Please rename your condition so that the name contains exactly one of "
+ "'+' or '-'.\n"
+ f"Got: '{label}'"
+ )
def define_model(name: str,
@@ -70,9 +141,17 @@ def define_model(name: str,
params : LinearParams
A dataclass containing k_pre, m_pre, k_sel, and m_sel.
"""
- with pyro.plate(f"{name}_condition_parameters", data.num_condition_rep):
+ if priors.m_is_selection is None:
+ m_scale_arr = jnp.full(data.num_condition_rep, priors.m_scale)
+ else:
+ m_scale_arr = jnp.array([
+ priors.m_scale_plus if sel else priors.m_scale_minus
+ for sel in priors.m_is_selection
+ ])
+
+ with pyro.plate(f"{name}_condition_parameters", data.num_condition_rep) as idx:
growth_k = pyro.sample(f"{name}_k", dist.Normal(priors.k_loc, priors.k_scale))
- growth_m = pyro.sample(f"{name}_m", dist.Normal(priors.m_loc, priors.m_scale))
+ growth_m = pyro.sample(f"{name}_m", dist.Normal(priors.m_loc, m_scale_arr[idx]))
k_pre = growth_k[data.map_condition_pre]
m_pre = growth_m[data.map_condition_pre]
@@ -151,10 +230,12 @@ def get_hyperparameters():
Get default values for the model hyperparameters.
"""
parameters = {}
- parameters["k_loc"] = 0.025
+ parameters["k_loc"] = 0.020
parameters["k_scale"] = 0.1
parameters["m_loc"] = 0.0
- parameters["m_scale"] = 0.05
+ parameters["m_scale"] = 0.01 # used when m_is_selection is None
+ parameters["m_scale_minus"] = 0.001 # tight prior for '-' control conditions
+ parameters["m_scale_plus"] = 0.01 # normal prior for '+' selection conditions
return parameters
@@ -170,7 +251,7 @@ def get_guesses(name, data):
Falls back to hard-coded defaults when ``data.ln_cfu`` or ``data.t_sel``
are not 7-D tensors.
"""
- _DEFAULT_K = 0.025
+ _DEFAULT_K = 0.020
_DEFAULT_SCALE = 0.01
num_cond_rep = data.num_condition_rep
@@ -241,11 +322,43 @@ def get_guesses(name, data):
return guesses
-def get_priors():
- return ModelPriors(**get_hyperparameters())
+def get_priors(condition_labels: Optional[List[str]] = None) -> ModelPriors:
+ """
+ Build a :class:`ModelPriors` for the linear growth component.
+
+ When ``condition_labels`` is supplied (a list of condition names in the
+ same order as the model's ``num_condition_rep`` axis), each label is
+ parsed for '+' or '-' to set ``m_is_selection``. This enables the
+ tighter ``m_scale_minus`` prior for control conditions and the normal
+ ``m_scale_plus`` prior for selection conditions, breaking the
+ optimization degeneracy that otherwise causes systematic theta
+ undershoot.
+
+ When ``condition_labels`` is ``None`` (the default), ``m_is_selection``
+ is left as ``None`` and ``m_scale`` is used for all conditions
+ (backward-compatible behaviour).
+
+ Parameters
+ ----------
+ condition_labels : list of str, optional
+ Ordered condition names, one per ``condition_rep`` index.
+ Each name must contain exactly one of '+' or '-'.
+
+ Returns
+ -------
+ ModelPriors
+ """
+ hypers = get_hyperparameters()
+ if condition_labels is not None:
+ hypers["m_is_selection"] = tuple(
+ _parse_condition_label(label) for label in condition_labels
+ )
+ return ModelPriors(**hypers)
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth/power.py b/src/tfscreen/tfmodel/generative/components/growth/power.py
index e5b7180d..8eb3e57b 100644
--- a/src/tfscreen/tfmodel/generative/components/growth/power.py
+++ b/src/tfscreen/tfmodel/generative/components/growth/power.py
@@ -156,7 +156,7 @@ def get_hyperparameters():
Get default values for the model hyperparameters.
"""
parameters = {}
- parameters["k_loc"] = 0.025
+ parameters["k_loc"] = 0.020
parameters["k_scale"] = 0.1
parameters["m_loc"] = 0.0
parameters["m_scale"] = 0.01
@@ -174,7 +174,7 @@ def get_guesses(name, data):
_DEFAULT_SCALE = 0.01
guesses = {}
- guesses[f"{name}_k_locs"] = jnp.full(num_cond_rep, 0.025, dtype=float)
+ guesses[f"{name}_k_locs"] = jnp.full(num_cond_rep, 0.020, dtype=float)
guesses[f"{name}_k_scales"] = jnp.full(num_cond_rep, _DEFAULT_SCALE, dtype=float)
guesses[f"{name}_m_locs"] = jnp.zeros(num_cond_rep, dtype=float)
guesses[f"{name}_m_scales"] = jnp.full(num_cond_rep, _DEFAULT_SCALE, dtype=float)
@@ -189,6 +189,8 @@ def get_priors():
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth/saturation.py b/src/tfscreen/tfmodel/generative/components/growth/saturation.py
index a5532e56..393cade4 100644
--- a/src/tfscreen/tfmodel/generative/components/growth/saturation.py
+++ b/src/tfscreen/tfmodel/generative/components/growth/saturation.py
@@ -128,9 +128,9 @@ def get_hyperparameters():
Get default values for the model hyperparameters.
"""
parameters = {}
- parameters["min_loc"] = 0.025
+ parameters["min_loc"] = 0.020
parameters["min_scale"] = 0.1
- parameters["max_loc"] = 0.025
+ parameters["max_loc"] = 0.030
parameters["max_scale"] = 0.1
return parameters
@@ -144,9 +144,9 @@ def get_guesses(name, data):
_DEFAULT_SCALE = 0.01
guesses = {}
- guesses[f"{name}_min_locs"] = jnp.full(num_cond_rep, 0.025, dtype=float)
+ guesses[f"{name}_min_locs"] = jnp.full(num_cond_rep, 0.020, dtype=float)
guesses[f"{name}_min_scales"] = jnp.full(num_cond_rep, _DEFAULT_SCALE, dtype=float)
- guesses[f"{name}_max_locs"] = jnp.full(num_cond_rep, 0.025, dtype=float)
+ guesses[f"{name}_max_locs"] = jnp.full(num_cond_rep, 0.030, dtype=float)
guesses[f"{name}_max_scales"] = jnp.full(num_cond_rep, _DEFAULT_SCALE, dtype=float)
return guesses
@@ -157,6 +157,8 @@ def get_priors():
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi.py b/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi.py
index 84b391de..86099bfa 100644
--- a/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi.py
+++ b/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi.py
@@ -160,6 +160,8 @@ def get_priors() -> ModelPriors:
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_k.py b/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_k.py
index 81227f69..f405751a 100644
--- a/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_k.py
+++ b/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_k.py
@@ -198,6 +198,8 @@ def get_priors() -> ModelPriors:
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_tau.py b/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_tau.py
index 7ad8017a..edc64c9f 100644
--- a/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_tau.py
+++ b/src/tfscreen/tfmodel/generative/components/growth_transition/baranyi_tau.py
@@ -198,6 +198,8 @@ def get_priors() -> ModelPriors:
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth_transition/memory.py b/src/tfscreen/tfmodel/generative/components/growth_transition/memory.py
index 6c42f789..7cba807a 100644
--- a/src/tfscreen/tfmodel/generative/components/growth_transition/memory.py
+++ b/src/tfscreen/tfmodel/generative/components/growth_transition/memory.py
@@ -194,6 +194,8 @@ def get_priors() -> ModelPriors:
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/growth_transition/two_pop.py b/src/tfscreen/tfmodel/generative/components/growth_transition/two_pop.py
index e9da6adc..8b0e15aa 100644
--- a/src/tfscreen/tfmodel/generative/components/growth_transition/two_pop.py
+++ b/src/tfscreen/tfmodel/generative/components/growth_transition/two_pop.py
@@ -235,6 +235,8 @@ def get_priors() -> ModelPriors:
def get_extract_specs(ctx):
+ if "condition_rep" not in ctx.growth_tm.map_groups:
+ return []
cond_rep_cols = (["condition_rep"] if ctx.growth_shares_replicates
else ["replicate", "condition_rep"])
return [dict(
diff --git a/src/tfscreen/tfmodel/generative/components/theta/_simple.py b/src/tfscreen/tfmodel/generative/components/theta/_simple.py
index 5a804438..37940331 100644
--- a/src/tfscreen/tfmodel/generative/components/theta/_simple.py
+++ b/src/tfscreen/tfmodel/generative/components/theta/_simple.py
@@ -4,8 +4,8 @@
This component is intended for the *calibration* / pre-fit pass of the
two-stage hierarchical workflow. It treats the per-condition fractional
occupancy ``theta`` as a fully known, deterministic quantity supplied by
-the caller — typically computed from a thermodynamic model evaluated at
-the wildtype TF parameters.
+the caller — typically the per-genotype observed binding values from the
+calibration dataset.
There are no sampled latent variables in this component:
- ``define_model`` registers zero ``pyro.sample`` sites and only emits a
@@ -13,19 +13,26 @@
- ``guide`` is the same pure no-op (no ``pyro.param`` and no
``pyro.sample`` sites), keeping model/guide sample sets symmetric.
-Because ``theta`` is identical across all genotypes, the broadcast tensor
-returned in ``ThetaParam.theta`` has shape
-``(num_titrant_name, num_titrant_conc, num_genotype)`` but is constant
-along the genotype axis.
+``theta_values`` in ``ModelPriors`` may have one of two shapes:
-The population moments ``(mu, sigma)`` are derived from the supplied
-``theta_values``:
-- ``mu`` = ``logit(theta_values)`` (with eps clipping)
-- ``sigma`` = a small floor (``priors.sigma_floor``)
+* ``(num_titrant_name, num_titrant_conc)`` — a single shared curve
+ broadcast identically to every genotype. Legacy / backward-compatible
+ path; rarely useful in practice.
-These match the contract used by downstream transformations (e.g.
-``transformation.logit_norm``) that need a population mean / scale even
-when theta is "known".
+* ``(num_titrant_name, num_titrant_conc, num_genotype)`` — per-genotype
+ theta values. The calibration pre-fit uses this path so that each
+ genotype's observed binding curve is used independently rather than
+ being replaced by a population average.
+
+The population moments ``(mu, sigma)`` are derived from ``theta_values``:
+
+* 2-D input: ``mu = logit(theta_values)[..., None]``;
+ ``sigma = priors.sigma_floor`` (scalar floor broadcast).
+* 3-D input: ``mu = mean(logit(theta_values), axis=-1, keepdims=True)``;
+ ``sigma = max(std(logit(theta_values), axis=-1, keepdims=True), floor)``.
+
+Both variants return moments of shape ``(T, C, 1)`` to satisfy the
+contract used by downstream transformations (e.g. ``transformation.logit_norm``).
"""
import jax.numpy as jnp
@@ -64,6 +71,8 @@ class ModelPriors:
theta_values: jnp.ndarray
sigma_floor: float = field(pytree_node=False, default=0.01)
+ # theta_values may be shape (T, C) or (T, C, G).
+ # See module docstring for semantics of each shape.
@dataclass(frozen=True)
@@ -89,9 +98,9 @@ class ThetaParam:
(passed through from ``data.titrant_conc``).
"""
- theta: jnp.ndarray
- mu: jnp.ndarray
- sigma: jnp.ndarray
+ theta: jnp.ndarray # (T, C, G) — per-genotype after construction
+ mu: jnp.ndarray # (T, C, 1) — population mean of logit(theta)
+ sigma: jnp.ndarray # (T, C, 1) — population scale of logit(theta)
concentrations: jnp.ndarray
@@ -106,19 +115,30 @@ def _build_theta_param(name: str,
difference is that ``define_model`` emits a ``pyro.deterministic``
site for inspection, while ``guide`` does not (sites in a guide must
not be observed).
- """
- theta_values = priors.theta_values # (num_titrant_name, num_titrant_conc)
+ Handles two ``priors.theta_values`` shapes:
- # Broadcast across genotype axis -> (Name, Conc, Genotype).
- theta = jnp.broadcast_to(
- theta_values[..., None],
- (data.num_titrant_name, data.num_titrant_conc, data.num_genotype),
- )
-
- # Population moments in logit-space, shape (Name, Conc, 1)
- mu = _logit(theta_values)[..., None]
- sigma = jnp.full_like(mu, priors.sigma_floor)
+ * ``(T, C)`` — broadcast identically to all G genotypes.
+ * ``(T, C, G)`` — per-genotype; used as-is (no broadcast).
+ """
+ theta_values = priors.theta_values
+
+ if theta_values.ndim == 2:
+ # Legacy / shared-curve path: broadcast (T, C) → (T, C, G).
+ theta = jnp.broadcast_to(
+ theta_values[..., None],
+ (data.num_titrant_name, data.num_titrant_conc, data.num_genotype),
+ )
+ # Population moments from the shared curve.
+ mu = _logit(theta_values)[..., None] # (T, C, 1)
+ sigma = jnp.full_like(mu, priors.sigma_floor)
+ else:
+ # Per-genotype path: theta_values already (T, C, G).
+ theta = theta_values
+ logit_tv = _logit(theta_values) # (T, C, G)
+ mu = jnp.mean(logit_tv, axis=-1, keepdims=True) # (T, C, 1)
+ std = jnp.std(logit_tv, axis=-1, keepdims=True) # (T, C, 1)
+ sigma = jnp.maximum(std, priors.sigma_floor)
if register_deterministic:
pyro.deterministic(f"{name}_theta", theta)
diff --git a/src/tfscreen/tfmodel/generative/components/theta/hill_geno.py b/src/tfscreen/tfmodel/generative/components/theta/hill_geno.py
index dd9a5ee0..3fef9861 100644
--- a/src/tfscreen/tfmodel/generative/components/theta/hill_geno.py
+++ b/src/tfscreen/tfmodel/generative/components/theta/hill_geno.py
@@ -4,7 +4,7 @@
import numpyro as pyro
import numpyro.distributions as dist
from flax.struct import dataclass
-from typing import Dict, Any
+from typing import Dict, Any, Optional
from tfscreen.tfmodel.data_class import DataClass
@@ -74,8 +74,63 @@ class ThetaParam:
theta_high: jnp.ndarray
log_hill_K: jnp.ndarray
hill_n: jnp.ndarray
- mu: jnp.ndarray
- sigma: jnp.ndarray
+ mu: Optional[jnp.ndarray]
+ sigma: Optional[jnp.ndarray]
+
+
+@dataclass(frozen=True)
+class SimPriors:
+ """
+ Parameters controlling perturbation-based theta simulation for hill_geno.
+
+ Genotypes are drawn from one of four phenotype categories:
+
+ * **normal** (probability ``1 - p_stuck_bound - p_never_binds - p_inverted``):
+ Hill curve perturbed around the wildtype reference.
+ * **stuck-bound** (``p_stuck_bound``): theta ≈ ``wt_theta_low`` throughout.
+ * **never-binds** (``p_never_binds``): theta ≈ ``wt_theta_high`` throughout.
+ * **inverted** (``p_inverted``): curve runs from ``wt_theta_high`` to ``wt_theta_low``.
+
+ Attributes
+ ----------
+ wt_theta_low : float
+ Wildtype theta at zero ligand concentration.
+ wt_theta_high : float
+ Wildtype theta at saturating ligand concentration.
+ wt_log_K : float
+ Wildtype log(K_D); sets the Hill-curve midpoint.
+ wt_hill_n : float
+ Wildtype Hill coefficient (cooperativity).
+ sigma_logit_low : float
+ Std dev of additive Normal noise on logit(theta_low) for normal genotypes.
+ sigma_logit_delta : float
+ Std dev of additive Normal noise on logit_delta (= logit_high - logit_low)
+ for normal and inverted genotypes.
+ sigma_log_K : float
+ Std dev of additive Normal noise on log_K for all genotypes.
+ sigma_log_n : float
+ Std dev of additive Normal noise on log(hill_n) for all genotypes.
+ p_stuck_bound : float
+ Fraction of genotypes assigned to the stuck-bound category.
+ p_never_binds : float
+ Fraction of genotypes assigned to the never-binds category.
+ p_inverted : float
+ Fraction of genotypes assigned to the inverted category.
+ """
+
+ wt_theta_low: float
+ wt_theta_high: float
+ wt_log_K: float
+ wt_hill_n: float
+
+ sigma_logit_low: float
+ sigma_logit_delta: float
+ sigma_log_K: float
+ sigma_log_n: float
+
+ p_stuck_bound: float
+ p_never_binds: float
+ p_inverted: float
def _population_moments(
@@ -362,6 +417,126 @@ def guide(name: str,
mu=mu, sigma=sigma)
+def simulate(name: str,
+ data: DataClass,
+ sim_priors: SimPriors,
+ rng_key) -> tuple:
+ """
+ Perturbation-based theta simulation for hill_geno.
+
+ Generates per-genotype Hill curves by perturbing a wildtype reference.
+ Genotypes are assigned to one of four phenotype categories (normal,
+ stuck-bound, never-binds, inverted) and their parameters drawn from
+ Normal distributions centered on the wildtype reference.
+
+ The returned ``ThetaParam`` has ``mu=None`` and ``sigma=None`` because
+ population-moment hyperparameters are not sampled in this path. Downstream
+ ``run_model`` and ``_theta_param_to_df`` only read the per-genotype fields
+ (``theta_low``, ``theta_high``, ``log_hill_K``, ``hill_n``), so this is safe.
+
+ Parameters
+ ----------
+ name : str
+ Unused; present for interface consistency with ``define_model``.
+ data : DataClass
+ Must expose ``num_genotype`` and ``log_titrant_conc``.
+ sim_priors : SimPriors
+ Wildtype reference and perturbation distributions.
+ rng_key : jax.random.PRNGKey
+ Seed for NumPy RNG (converted internally).
+
+ Returns
+ -------
+ theta_gc : np.ndarray, shape (num_genotype, num_titrant_conc)
+ Fractional occupancy for each genotype at each unique concentration.
+ theta_param : ThetaParam
+ Per-genotype fields shape ``(1, num_genotype)``; ``mu`` and ``sigma``
+ are ``None``.
+ """
+ _eps = 1e-6
+ seed = int(jax.random.randint(rng_key, shape=(), minval=0, maxval=2**30))
+ rng = np.random.default_rng(seed)
+
+ G = data.num_genotype
+ log_conc = np.array(data.log_titrant_conc) # (C,)
+
+ # Wildtype reference in logit space
+ wt_logit_low = float(np.log(np.clip(sim_priors.wt_theta_low, _eps, 1 - _eps)
+ / (1 - np.clip(sim_priors.wt_theta_low, _eps, 1 - _eps))))
+ wt_logit_high = float(np.log(np.clip(sim_priors.wt_theta_high, _eps, 1 - _eps)
+ / (1 - np.clip(sim_priors.wt_theta_high, _eps, 1 - _eps))))
+ wt_logit_delta = wt_logit_high - wt_logit_low
+ wt_log_n = float(np.log(sim_priors.wt_hill_n))
+
+ # Assign phenotype categories
+ p_normal = 1.0 - sim_priors.p_stuck_bound - sim_priors.p_never_binds - sim_priors.p_inverted
+ if p_normal < 0:
+ raise ValueError(
+ "SimPriors mixture probabilities (p_stuck_bound + p_never_binds + p_inverted) "
+ "sum to more than 1.0"
+ )
+ probs = [p_normal, sim_priors.p_stuck_bound, sim_priors.p_never_binds, sim_priors.p_inverted]
+ categories = rng.choice(4, size=G, p=probs)
+
+ # Per-genotype parameters, initialised to wildtype
+ logit_low = np.full(G, wt_logit_low)
+ logit_delta = np.full(G, wt_logit_delta)
+ log_K = np.full(G, sim_priors.wt_log_K)
+ log_n = np.full(G, wt_log_n)
+
+ # Category 0: normal (wildtype-like, full perturbation on all parameters)
+ m0 = categories == 0
+ n0 = int(m0.sum())
+ logit_low[m0] += rng.normal(0.0, sim_priors.sigma_logit_low, n0)
+ logit_delta[m0] += rng.normal(0.0, sim_priors.sigma_logit_delta, n0)
+ log_K[m0] += rng.normal(0.0, sim_priors.sigma_log_K, n0)
+ log_n[m0] += rng.normal(0.0, sim_priors.sigma_log_n, n0)
+
+ # Category 1: stuck-bound (theta ≈ wt_theta_low, near-zero transition)
+ m1 = categories == 1
+ n1 = int(m1.sum())
+ logit_low[m1] = wt_logit_low + rng.normal(0.0, sim_priors.sigma_logit_low, n1)
+ logit_delta[m1] = rng.normal(0.0, 0.1, n1)
+ log_K[m1] += rng.normal(0.0, sim_priors.sigma_log_K, n1)
+ log_n[m1] += rng.normal(0.0, sim_priors.sigma_log_n, n1)
+
+ # Category 2: never-binds (theta ≈ wt_theta_high, near-zero transition)
+ m2 = categories == 2
+ n2 = int(m2.sum())
+ logit_low[m2] = wt_logit_high + rng.normal(0.0, sim_priors.sigma_logit_low, n2)
+ logit_delta[m2] = rng.normal(0.0, 0.1, n2)
+ log_K[m2] += rng.normal(0.0, sim_priors.sigma_log_K, n2)
+ log_n[m2] += rng.normal(0.0, sim_priors.sigma_log_n, n2)
+
+ # Category 3: inverted (theta goes from wt_theta_high to wt_theta_low)
+ m3 = categories == 3
+ n3 = int(m3.sum())
+ logit_low[m3] = wt_logit_high + rng.normal(0.0, sim_priors.sigma_logit_low, n3)
+ logit_delta[m3] = -wt_logit_delta + rng.normal(0.0, sim_priors.sigma_logit_delta, n3)
+ log_K[m3] += rng.normal(0.0, sim_priors.sigma_log_K, n3)
+ log_n[m3] += rng.normal(0.0, sim_priors.sigma_log_n, n3)
+
+ # Convert to probability space
+ theta_low_arr = 1.0 / (1.0 + np.exp(-logit_low)) # (G,)
+ theta_high_arr = 1.0 / (1.0 + np.exp(-(logit_low + logit_delta))) # (G,)
+ hill_n_arr = np.exp(log_n) # (G,)
+
+ # Compute theta: (G, C)
+ occupancy = 1.0 / (1.0 + np.exp(-hill_n_arr[:, None] * (log_conc[None, :] - log_K[:, None])))
+ theta_gc = theta_low_arr[:, None] + (theta_high_arr - theta_low_arr)[:, None] * occupancy
+
+ theta_param = ThetaParam(
+ theta_low = jnp.array(theta_low_arr)[None, :], # (1, G)
+ theta_high = jnp.array(theta_high_arr)[None, :], # (1, G)
+ log_hill_K = jnp.array(log_K)[None, :], # (1, G)
+ hill_n = jnp.array(hill_n_arr)[None, :], # (1, G)
+ mu = None,
+ sigma = None,
+ )
+
+ return theta_gc, theta_param
+
+
def run_model(theta_param: ThetaParam, data: DataClass) -> jnp.ndarray:
"""
Calculates fractional occupancy (theta) using the Hill equation.
@@ -439,6 +614,33 @@ def get_hyperparameters() -> Dict[str, Any]:
return parameters
+def get_sim_hyperparameters() -> Dict[str, Any]:
+ """
+ Default hyperparameters for perturbation-based simulation (``SimPriors``).
+
+ Wildtype reference is a full-range decreasing Hill curve (theta: ~1→~0)
+ matching the lac/IPTG system. Perturbation sigmas allow realistic spread
+ across the genotype library.
+
+ Returns
+ -------
+ dict[str, Any]
+ """
+ return {
+ "wt_theta_low": 0.99,
+ "wt_theta_high": 0.01,
+ "wt_log_K": -4.1, # ln(0.017 mM) — lac/IPTG default
+ "wt_hill_n": 2.0,
+ "sigma_logit_low": 0.5,
+ "sigma_logit_delta": 0.5,
+ "sigma_log_K": 0.5,
+ "sigma_log_n": 0.3,
+ "p_stuck_bound": 0.05,
+ "p_never_binds": 0.05,
+ "p_inverted": 0.02,
+ }
+
+
def get_guesses(name: str, data: DataClass) -> Dict[str, Any]:
"""
Gets initial guess values for model parameters.
diff --git a/src/tfscreen/tfmodel/generative/components/theta/hill_mut.py b/src/tfscreen/tfmodel/generative/components/theta/hill_mut.py
index 35bf0272..d3b66b51 100644
--- a/src/tfscreen/tfmodel/generative/components/theta/hill_mut.py
+++ b/src/tfscreen/tfmodel/generative/components/theta/hill_mut.py
@@ -32,10 +32,8 @@
import numpyro.distributions as dist
import pandas as pd
from flax.struct import dataclass
-from typing import Dict, Any
-
+from typing import Dict, Any, Optional, Union
from functools import partial
-from typing import Union
from tfscreen.tfmodel.data_class import GrowthData, BindingData
from tfscreen.genetics.build_mut_geno_matrix import apply_pair_matrix, apply_mut_matrix
@@ -78,8 +76,68 @@ class ThetaParam:
theta_high: jnp.ndarray
log_hill_K: jnp.ndarray
hill_n: jnp.ndarray
- mu: jnp.ndarray
- sigma: jnp.ndarray
+ mu: Optional[jnp.ndarray]
+ sigma: Optional[jnp.ndarray]
+
+
+@dataclass(frozen=True)
+class SimPriors:
+ """
+ Parameters controlling perturbation-based theta simulation for hill_mut.
+
+ Mutations perturb each Hill parameter by additive Normal noise in
+ transformed space. The wildtype genotype (no mutations) receives exactly
+ the reference curve. Per-genotype values are assembled by summing per-mutation
+ deltas over each genotype's mutation set, optionally plus pairwise epistasis.
+
+ Attributes
+ ----------
+ wt_theta_low : float
+ Wildtype theta at zero ligand concentration.
+ wt_theta_high : float
+ Wildtype theta at saturating ligand concentration.
+ wt_log_K : float
+ Wildtype log(K_D); sets the Hill-curve midpoint.
+ wt_hill_n : float
+ Wildtype Hill coefficient (cooperativity).
+ sigma_d_logit_low : float
+ Std dev of per-mutation additive delta on logit(theta_low).
+ sigma_d_logit_delta : float
+ Std dev of per-mutation additive delta on logit_delta
+ (= logit_high - logit_low).
+ sigma_d_log_K : float
+ Std dev of per-mutation additive delta on log_K. Primary binding-
+ affinity effect; typically the largest perturbation.
+ sigma_d_log_n : float
+ Std dev of per-mutation additive delta on log(hill_n).
+ epi_tau_scale : float
+ HalfCauchy scale for the global sparsity parameter τ of the
+ regularized horseshoe prior on pairwise epistasis. Setting this to
+ ``0.0`` makes τ=0 exactly, disabling all epistasis. Ignored when
+ ``data.num_pair == 0``.
+ epi_slab_scale : float
+ Typical magnitude of a large (non-shrunk) epistasis effect; sets the
+ scale of the InvGamma slab (c²).
+ epi_slab_df : float
+ Degrees of freedom for the InvGamma slab prior on c². Usually 4.
+ """
+
+ wt_theta_low: float
+ wt_theta_high: float
+ wt_log_K: float
+ wt_hill_n: float
+
+ sigma_d_logit_low: float
+ sigma_d_logit_delta: float
+ sigma_d_log_K: float
+ sigma_d_log_n: float
+
+ # Regularized horseshoe prior parameters for pairwise epistasis.
+ # Used only when ``data.num_pair > 0``.
+ # Set ``epi_tau_scale=0.0`` to disable epistasis exactly (τ=0 → zero effects).
+ epi_tau_scale: float # HalfCauchy scale for global τ
+ epi_slab_scale: float # typical magnitude of a large epistasis effect
+ epi_slab_df: float # InvGamma degrees of freedom (slab; usually 4)
# ---------------------------------------------------------------------------
@@ -216,16 +274,23 @@ def define_model(name: str,
# ------------------------------------------------------------------
# Mutation delta offsets: shape (T, num_mutation)
# ------------------------------------------------------------------
- with pyro.plate(f"{name}_titrant_mut_outer_plate", T, dim=-2):
- with pyro.plate(f"{name}_mutation_plate", num_mut, dim=-1):
- d_low_off = pyro.sample(
- f"{name}_d_logit_low_offset", dist.Normal(0.0, 1.0))
- d_delta_off = pyro.sample(
- f"{name}_d_logit_delta_offset", dist.Normal(0.0, 1.0))
- d_K_off = pyro.sample(
- f"{name}_d_log_hill_K_offset", dist.Normal(0.0, 1.0))
- d_n_off = pyro.sample(
- f"{name}_d_log_hill_n_offset", dist.Normal(0.0, 1.0))
+ if num_mut > 0:
+ with pyro.plate(f"{name}_titrant_mut_outer_plate", T, dim=-2):
+ with pyro.plate(f"{name}_mutation_plate", num_mut, dim=-1):
+ d_low_off = pyro.sample(
+ f"{name}_d_logit_low_offset", dist.Normal(0.0, 1.0))
+ d_delta_off = pyro.sample(
+ f"{name}_d_logit_delta_offset", dist.Normal(0.0, 1.0))
+ d_K_off = pyro.sample(
+ f"{name}_d_log_hill_K_offset", dist.Normal(0.0, 1.0))
+ d_n_off = pyro.sample(
+ f"{name}_d_log_hill_n_offset", dist.Normal(0.0, 1.0))
+ else:
+ # num_mut == 0 (wt-only library, e.g. stratified pool building):
+ # numpyro.plate requires size > 0, so skip the plate and return
+ # empty offset arrays. apply_mut_matrix with empty COO indices
+ # produces zeros, so per-genotype params equal the WT values.
+ d_low_off = d_delta_off = d_K_off = d_n_off = jnp.zeros((T, 0))
# ------------------------------------------------------------------
# Optional epistasis: shape (T, num_pair)
@@ -571,6 +636,142 @@ def _lam_tilde(lam):
sigma=sigma)
+# ---------------------------------------------------------------------------
+# Perturbation-based simulation
+# ---------------------------------------------------------------------------
+
+def simulate(name: str,
+ data,
+ sim_priors: SimPriors,
+ rng_key) -> tuple:
+ """
+ Perturbation-based theta simulation for hill_mut.
+
+ Generates per-genotype Hill curves by assigning a wildtype reference curve
+ and adding per-mutation additive deltas in transformed parameter space.
+ The wildtype genotype (no mutations present) receives the reference curve
+ exactly. Pairwise epistasis is sampled and applied when ``data.num_pair > 0``.
+
+ Pairwise epistasis is sampled from a regularized horseshoe prior when
+ ``data.num_pair > 0``. Setting ``sim_priors.epi_tau_scale=0.0`` makes
+ the global scale τ=0 exactly, which forces all epistasis effects to zero
+ (a clean "no epistasis" toggle).
+
+ The returned ``ThetaParam`` has ``mu=None`` and ``sigma=None``; downstream
+ ``run_model`` only reads the per-genotype fields.
+
+ Parameters
+ ----------
+ name : str
+ Unused; present for interface consistency with ``define_model``.
+ data : SimData or compatible
+ Must expose ``num_genotype``, ``num_mutation``, ``num_pair``,
+ ``log_titrant_conc``, ``mut_nnz_mut_idx``, ``mut_nnz_geno_idx``,
+ and (when ``num_pair > 0``) ``pair_nnz_pair_idx``,
+ ``pair_nnz_geno_idx``.
+ sim_priors : SimPriors
+ Wildtype reference and perturbation distributions.
+ rng_key : jax.random.PRNGKey
+ Seed for NumPy RNG (converted internally).
+
+ Returns
+ -------
+ theta_gc : np.ndarray, shape (num_genotype, num_titrant_conc)
+ Fractional occupancy for each genotype at each unique concentration.
+ theta_param : ThetaParam
+ Per-genotype fields shape ``(1, num_genotype)``; ``mu`` and ``sigma``
+ are ``None``.
+ """
+ _eps = 1e-6
+ seed = int(jax.random.randint(rng_key, shape=(), minval=0, maxval=2**30))
+ rng = np.random.default_rng(seed)
+
+ G = data.num_genotype
+ M = data.num_mutation
+ log_conc = np.array(data.log_titrant_conc) # (C,)
+
+ # Scatter callables: accept (T, M) or (T, P), return (T, G)
+ mut_scatter = partial(apply_mut_matrix,
+ mut_nnz_mut_idx=jnp.array(data.mut_nnz_mut_idx),
+ mut_nnz_geno_idx=jnp.array(data.mut_nnz_geno_idx),
+ num_genotype=G)
+
+ # Wildtype reference in logit space
+ wt_logit_low = float(np.log(np.clip(sim_priors.wt_theta_low, _eps, 1 - _eps)
+ / (1 - np.clip(sim_priors.wt_theta_low, _eps, 1 - _eps))))
+ wt_logit_high = float(np.log(np.clip(sim_priors.wt_theta_high, _eps, 1 - _eps)
+ / (1 - np.clip(sim_priors.wt_theta_high, _eps, 1 - _eps))))
+ wt_logit_delta = wt_logit_high - wt_logit_low
+ wt_log_n = float(np.log(sim_priors.wt_hill_n))
+
+ def _scatter_mut(d_1d):
+ """Scatter (M,) deltas to (G,) via the mutation-genotype COO map."""
+ return np.array(mut_scatter(jnp.array(d_1d)[None, :]))[0]
+
+ # Sample per-mutation deltas and assemble per-genotype parameters: (G,)
+ logit_low = wt_logit_low + _scatter_mut(rng.normal(0.0, sim_priors.sigma_d_logit_low, M))
+ logit_delta = wt_logit_delta + _scatter_mut(rng.normal(0.0, sim_priors.sigma_d_logit_delta, M))
+ log_K = sim_priors.wt_log_K + _scatter_mut(rng.normal(0.0, sim_priors.sigma_d_log_K, M))
+ log_n = wt_log_n + _scatter_mut(rng.normal(0.0, sim_priors.sigma_d_log_n, M))
+
+ # Optional pairwise epistasis — regularized horseshoe prior
+ # epi_tau_scale=0.0 → tau=0 exactly → all epistasis effects are zero.
+ if data.num_pair > 0:
+ pair_scatter = partial(apply_pair_matrix,
+ pair_nnz_pair_idx=jnp.array(data.pair_nnz_pair_idx),
+ pair_nnz_geno_idx=jnp.array(data.pair_nnz_geno_idx),
+ num_genotype=G)
+ P = data.num_pair
+
+ def _scatter_pair(e_1d):
+ """Scatter (P,) epistasis to (G,) via the pair-genotype COO map."""
+ return np.array(pair_scatter(jnp.array(e_1d)[None, :]))[0]
+
+ # Global sparsity scale τ ~ HalfCauchy(epi_tau_scale).
+ # When epi_tau_scale == 0.0, tau == 0 and all effects are exactly zero.
+ tau = float(np.abs(rng.standard_cauchy())) * sim_priors.epi_tau_scale
+
+ if tau > 0.0:
+ # Slab variance c² ~ InvGamma(slab_df/2, slab_df*slab_scale²/2).
+ # Sampled as 1/Gamma(α, scale=1/β) with α=slab_df/2, β=slab_df*slab_scale²/2.
+ c2 = 1.0 / rng.gamma(
+ shape=sim_priors.epi_slab_df / 2.0,
+ scale=2.0 / (sim_priors.epi_slab_df * sim_priors.epi_slab_scale ** 2))
+
+ def _horseshoe(size):
+ """Draw horseshoe-shrunk epistasis effects of shape (size,)."""
+ lam = np.abs(rng.standard_cauchy(size)) # per-pair local scale
+ lam_tilde = np.sqrt(c2 * lam**2
+ / (c2 + tau**2 * lam**2)) # slab-regularised scale
+ return rng.standard_normal(size) * tau * lam_tilde
+
+ logit_low += _scatter_pair(_horseshoe(P))
+ logit_delta += _scatter_pair(_horseshoe(P))
+ log_K += _scatter_pair(_horseshoe(P))
+ log_n += _scatter_pair(_horseshoe(P))
+ # else: tau == 0.0 → all epistasis effects are exactly zero; nothing to add
+
+ # Convert to probability space
+ theta_low_arr = 1.0 / (1.0 + np.exp(-logit_low)) # (G,)
+ theta_high_arr = 1.0 / (1.0 + np.exp(-(logit_low + logit_delta))) # (G,)
+ hill_n_arr = np.exp(log_n) # (G,)
+
+ # Compute theta: (G, C)
+ occupancy = 1.0 / (1.0 + np.exp(-hill_n_arr[:, None] * (log_conc[None, :] - log_K[:, None])))
+ theta_gc = theta_low_arr[:, None] + (theta_high_arr - theta_low_arr)[:, None] * occupancy
+
+ theta_param = ThetaParam(
+ theta_low = jnp.array(theta_low_arr)[None, :], # (1, G)
+ theta_high = jnp.array(theta_high_arr)[None, :],
+ log_hill_K = jnp.array(log_K)[None, :],
+ hill_n = jnp.array(hill_n_arr)[None, :],
+ mu = None,
+ sigma = None,
+ )
+
+ return theta_gc, theta_param
+
+
# ---------------------------------------------------------------------------
# run_model – identical to hill.py
# ---------------------------------------------------------------------------
@@ -608,16 +809,16 @@ def get_population_moments(theta_param: ThetaParam, data) -> tuple:
def get_hyperparameters() -> Dict[str, Any]:
p = {}
- p["theta_logit_low_wt_loc"] = 2.0
+ p["theta_logit_low_wt_loc"] = 5.0
p["theta_logit_low_wt_scale"] = 2.0
- p["theta_logit_delta_wt_loc"] = -4.0
+ p["theta_logit_delta_wt_loc"] = -9.0
p["theta_logit_delta_wt_scale"] = 2.0
p["theta_log_hill_K_wt_loc"] = -4.1
p["theta_log_hill_K_wt_scale"] = 2.0
p["theta_log_hill_n_wt_loc"] = 0.7
p["theta_log_hill_n_wt_scale"] = 0.3
- p["theta_sigma_d_logit_low_scale"] = 1.0
- p["theta_sigma_d_logit_delta_scale"] = 1.0
+ p["theta_sigma_d_logit_low_scale"] = 5.0
+ p["theta_sigma_d_logit_delta_scale"] = 5.0
p["theta_sigma_d_log_hill_K_scale"] = 1.0
p["theta_sigma_d_log_hill_n_scale"] = 0.5
p["theta_epi_tau_scale"] = 0.1
@@ -626,16 +827,44 @@ def get_hyperparameters() -> Dict[str, Any]:
return p
+def get_sim_hyperparameters() -> Dict[str, Any]:
+ """
+ Default hyperparameters for perturbation-based simulation (``SimPriors``).
+
+ The wildtype reference is a full-range decreasing Hill curve matching the
+ lac/IPTG system. Per-mutation effects primarily shift K_D (``sigma_d_log_K``)
+ with smaller effects on cooperativity and the baseline/ceiling occupancy.
+ Epistasis is small by default.
+
+ Returns
+ -------
+ dict[str, Any]
+ """
+ return {
+ "wt_theta_low": 0.99,
+ "wt_theta_high": 0.01,
+ "wt_log_K": -4.1, # ln(0.017 mM) — lac/IPTG default
+ "wt_hill_n": 2.0,
+ "sigma_d_logit_low": 0.3,
+ "sigma_d_logit_delta": 0.5,
+ "sigma_d_log_K": 0.5,
+ "sigma_d_log_n": 0.3,
+ "epi_tau_scale": 0.1, # matches theta_epi_tau_scale in get_hyperparameters()
+ "epi_slab_scale": 2.0, # matches theta_epi_slab_scale
+ "epi_slab_df": 4.0, # matches theta_epi_slab_df
+ }
+
+
def get_guesses(name: str, data: GrowthData) -> Dict[str, Any]:
T = data.num_titrant_name
M = data.num_mutation
guesses = {}
- guesses[f"{name}_logit_low_wt"] = jnp.full(T, 2.0)
- guesses[f"{name}_logit_delta_wt"] = jnp.full(T, -4.0)
+ guesses[f"{name}_logit_low_wt"] = jnp.full(T, 5.0)
+ guesses[f"{name}_logit_delta_wt"] = jnp.full(T, -9.0)
guesses[f"{name}_log_hill_K_wt"] = jnp.full(T, -4.1)
guesses[f"{name}_log_hill_n_wt"] = jnp.full(T, 0.7)
guesses[f"{name}_sigma_d_logit_low"] = jnp.full(T, 0.5)
- guesses[f"{name}_sigma_d_logit_delta"] = jnp.full(T, 0.5)
+ guesses[f"{name}_sigma_d_logit_delta"] = jnp.full(T, 4.0)
guesses[f"{name}_sigma_d_log_hill_K"] = jnp.full(T, 0.5)
guesses[f"{name}_sigma_d_log_hill_n"] = jnp.full(T, 0.3)
guesses[f"{name}_d_logit_low_offset"] = jnp.zeros((T, M))
diff --git a/src/tfscreen/tfmodel/generative/components/transformation/_congression.py b/src/tfscreen/tfmodel/generative/components/transformation/_congression.py
index 57609bc8..ac13f750 100644
--- a/src/tfscreen/tfmodel/generative/components/transformation/_congression.py
+++ b/src/tfscreen/tfmodel/generative/components/transformation/_congression.py
@@ -462,7 +462,7 @@ def get_guesses(name,data):
# Only add these if we are not in empirical mode.
# NOTE: get_guesses currently doesn't know the mode from the config easily
- # during initialization in model_class.py unless we pass it.
+ # during initialization in model_orchestrator.py unless we pass it.
# For now, we'll keep them as guesses; they just won't be used if not
# in the model/guide.
guesses[f"{name}_mu"] = 0.0
diff --git a/src/tfscreen/tfmodel/generative/model.py b/src/tfscreen/tfmodel/generative/model.py
index 9d8e2c70..10860140 100644
--- a/src/tfscreen/tfmodel/generative/model.py
+++ b/src/tfscreen/tfmodel/generative/model.py
@@ -7,6 +7,7 @@
import jax.numpy as jnp
import numpyro as pyro
+import numpyro.distributions as dist
def jax_model(data: DataClass,
priors: PriorsClass,
@@ -184,6 +185,31 @@ def jax_model(data: DataClass,
# real calculation
else:
+ # Pre-split (t = -t_pre) observations — direct constraint on ln_cfu0.
+ # ln_cfu0 shape: (num_rep, 1, num_cp, 1, 1, 1, batch_size)
+ # Squeeze broadcast dims to get (num_rep, num_cp, batch_size).
+ if getattr(data, "presplit", None) is not None:
+ ln_cfu0_3d = ln_cfu0[:, 0, :, 0, 0, 0, :]
+ ps = data.presplit
+ bi = data.growth.batch_idx
+ obs_t0 = ps.ln_cfu_t0[:, :, bi]
+ std_t0 = ps.ln_cfu_t0_std[:, :, bi]
+ mask_t0 = ps.good_mask[:, :, bi]
+ with pyro.plate("presplit_replicate",
+ size=ps.num_replicate, dim=-3):
+ with pyro.plate("presplit_condition_pre",
+ size=ps.num_condition_pre, dim=-2):
+ with pyro.plate("shared_genotype_plate",
+ size=data.growth.batch_size, dim=-1):
+ with pyro.handlers.scale(
+ scale=data.growth.scale_vector):
+ with pyro.handlers.mask(mask=mask_t0):
+ pyro.sample(
+ "presplit_obs",
+ dist.Normal(ln_cfu0_3d, std_t0),
+ obs=obs_t0,
+ )
+
# calculate observable (all tensors have correct dimensions)
g_pre, g_sel = calculate_growth(params=growth_params,
dk_geno=dk_geno,
diff --git a/src/tfscreen/tfmodel/generative/registry.py b/src/tfscreen/tfmodel/generative/registry.py
index df430dbc..ae58e6e2 100644
--- a/src/tfscreen/tfmodel/generative/registry.py
+++ b/src/tfscreen/tfmodel/generative/registry.py
@@ -21,6 +21,7 @@
from .components.activity import hierarchical_mut as activity_mut_decomp
from .components.activity import horseshoe_mut as activity_horseshoe_mut
+from .components.theta import _simple as theta_simple
from .components.theta import categorical_geno as theta_cat
from .components.theta import hill_geno as theta_hill
from .components.theta import hill_mut as theta_hill_mut
@@ -87,6 +88,7 @@
"logit": theta_rescale_logit,
},
"theta":{
+ "_simple":theta_simple,
"categorical_geno":theta_cat,
"hill_geno":theta_hill,
"hill_mut":theta_hill_mut,
diff --git a/src/tfscreen/tfmodel/inference/checkpoint_io.py b/src/tfscreen/tfmodel/inference/checkpoint_io.py
index a3fa4133..f9cc06e5 100644
--- a/src/tfscreen/tfmodel/inference/checkpoint_io.py
+++ b/src/tfscreen/tfmodel/inference/checkpoint_io.py
@@ -2,7 +2,7 @@
import dill
-def resolve_param_file(param_file, gm, out_prefix):
+def resolve_param_file(param_file, orchestrator, out_prefix):
"""
Return a posterior .h5 path, converting a MAP checkpoint on the fly if needed.
@@ -22,7 +22,7 @@ def resolve_param_file(param_file, gm, out_prefix):
----------
param_file : str
Path to a posterior .h5 file or a MAP checkpoint .pkl file.
- gm : ModelOrchestrator
+ orchestrator : ModelOrchestrator
Already-loaded model instance (from ``read_configuration``).
out_prefix : str
Output prefix for the calling script. The intermediate map posterior
@@ -52,7 +52,7 @@ def resolve_param_file(param_file, gm, out_prefix):
from tfscreen.tfmodel.inference.run_inference import RunInference
- ri = RunInference(gm, seed=0)
+ ri = RunInference(orchestrator, seed=0)
temp_svi = ri.setup_svi(guide_type="delta")
map_params = temp_svi.optim.get_params(chk_data["svi_state"].optim_state)
diff --git a/src/tfscreen/tfmodel/inference/posteriors.py b/src/tfscreen/tfmodel/inference/posteriors.py
index d5c2d468..d28c7c58 100644
--- a/src/tfscreen/tfmodel/inference/posteriors.py
+++ b/src/tfscreen/tfmodel/inference/posteriors.py
@@ -1,6 +1,28 @@
import numpy as np
import h5py
+# ---------------------------------------------------------------------------
+# Default quantile levels and column-name helper
+# ---------------------------------------------------------------------------
+
+_DEFAULT_QUANTILE_LEVELS = np.array([
+ 0.001, 0.005, 0.01, 0.025, 0.05,
+ 0.10, 0.159, 0.25, 0.50, 0.75, 0.841, 0.90,
+ 0.95, 0.975, 0.99, 0.995, 0.999,
+])
+
+
+def quantile_col(level):
+ """Return the standard column name for quantile *level*.
+
+ Uses ``f"q{level:g}"`` so trailing zeros are stripped:
+ ``quantile_col(0.5)`` → ``"q0.5"``,
+ ``quantile_col(0.025)`` → ``"q0.025"``,
+ ``quantile_col(0.159)`` → ``"q0.159"``.
+ """
+ return f"q{level:g}"
+
+
def get_posterior_samples(param_posteriors, param_name):
"""
Get posterior samples for a parameter, handling name fallbacks and HDF5.
@@ -39,7 +61,7 @@ def get_posterior_samples(param_posteriors, param_name):
keys_str = ", ".join(available_keys[:5]) + " ... " + ", ".join(available_keys[-5:])
else:
keys_str = ", ".join(available_keys)
-
+
error_msg = f"Parameter '{param_name}' not found in posteriors. Available keys: {keys_str}"
raise KeyError(error_msg)
@@ -49,25 +71,36 @@ def get_posterior_samples(param_posteriors, param_name):
def load_posteriors(posteriors, q_to_get=None):
"""
- Consolidates the reading logic for posterior samples and sets default quantiles.
+ Consolidate the reading logic for posterior samples and build quantile mapping.
Parameters
----------
posteriors : dict or str
- Assumes this is a dictionary of posteriors keying parameters to
- numpy arrays, a numpy.lib.npyio.NpzFile object, or a path to a
- .npz or .h5/.hdf5 file containing posterior samples for model
+ A dictionary of posteriors keying parameters to numpy arrays, a
+ ``numpy.lib.npyio.NpzFile`` object, or a path to a ``.npz`` or
+ ``.h5``/``.hdf5`` file containing posterior samples for model
parameters.
- q_to_get : dict, optional
- Dictionary mapping output column names to quantile values (between 0 and 1)
- to extract from the posterior samples. If None, a default set of quantiles
- is used (min, lower_95, lower_std, lower_quartile, median, upper_std,
- upper_quartile, upper_95, max).
+ q_to_get : array-like of float, optional
+ Quantile levels in [0, 1] to extract from posterior samples. Column
+ names are generated automatically as ``q{level:g}``
+ (e.g. ``q0.5`` for the median, ``q0.025`` for the 2.5th percentile).
+ If ``None``, a dense default set is used::
+
+ [0.001, 0.005, 0.01, 0.025, 0.05,
+ 0.10, 0.159, 0.25, 0.50, 0.75, 0.841, 0.90,
+ 0.95, 0.975, 0.99, 0.995, 0.999]
Returns
-------
tuple
- A tuple containing (q_to_get, param_posteriors).
+ ``(q_dict, param_posteriors)`` where ``q_dict`` maps column names to
+ quantile levels (e.g. ``{"q0.5": 0.5, "q0.025": 0.025, ...}``).
+
+ Raises
+ ------
+ ValueError
+ If ``q_to_get`` cannot be converted to a 1-D float array or any
+ value is outside [0, 1].
"""
# Load the posterior file
@@ -95,22 +128,24 @@ def load_posteriors(posteriors, q_to_get=None):
"path), an .h5 file (loaded or path), or an h5py group."
)
- # Named quantiles to pull from the posterior distribution
+ # Build the quantile level array
if q_to_get is None:
- q_to_get = {"min": 0.0,
- "lower_95": 0.025,
- "lower_std": 0.159,
- "lower_quartile": 0.25,
- "median": 0.5,
- "upper_quartile": 0.75,
- "upper_std": 0.841,
- "upper_95": 0.975,
- "max": 1.0}
-
- # make sure q_to_get is a dictionary
- if not isinstance(q_to_get, dict):
- raise ValueError(
- "q_to_get should be a dictionary keying column names to quantiles"
- )
-
- return q_to_get, param_posteriors
+ q_arr = _DEFAULT_QUANTILE_LEVELS
+ else:
+ try:
+ q_arr = np.atleast_1d(np.asarray(q_to_get, dtype=float)).ravel()
+ except (TypeError, ValueError):
+ raise ValueError(
+ "q_to_get should be a 1-D array-like of quantile levels in [0, 1]; "
+ f"got {type(q_to_get).__name__}"
+ )
+ if not np.all((q_arr >= 0) & (q_arr <= 1)):
+ raise ValueError(
+ "q_to_get values must all be in [0, 1]; "
+ f"got min={float(q_arr.min()):.4g}, max={float(q_arr.max()):.4g}"
+ )
+
+ # Build column-name → level dict used by all downstream extraction code
+ q_dict = {quantile_col(level): float(level) for level in q_arr}
+
+ return q_dict, param_posteriors
diff --git a/src/tfscreen/tfmodel/inference/run_inference.py b/src/tfscreen/tfmodel/inference/run_inference.py
index 85de3478..1a210f38 100644
--- a/src/tfscreen/tfmodel/inference/run_inference.py
+++ b/src/tfscreen/tfmodel/inference/run_inference.py
@@ -357,9 +357,14 @@ def scan_fn(data_on_gpu, carry, indices):
self._write_checkpoint(svi_state, out_prefix)
break
+ # Write a final checkpoint when the loop exits by reaching max_num_epochs
+ # (convergence already writes its own checkpoint via the break path above).
+ if not converged and total_steps > 0:
+ self._write_checkpoint(svi_state, out_prefix)
+
# Get final parameters
params = svi.get_params(svi_state)
-
+
return svi_state, params, converged
def _get_genotype_dim_map(self):
@@ -415,6 +420,16 @@ def _get_genotype_dim_map(self):
if in_plate:
continue
+ # Skip sites that belong to a non-genotype plate at the genotype
+ # dim — they are indexed by a different axis (e.g. condition_pre)
+ # and must NOT be treated as genotype-indexed.
+ in_other_plate_at_geno_dim = any(
+ frame.dim == genotype_dim
+ for frame in site.get("cond_indep_stack", [])
+ )
+ if in_other_plate_at_geno_dim:
+ continue
+
# Fallback for deterministics computed outside the plate
# but matching the genotype size at the expected dimension.
val = site["value"]
@@ -718,6 +733,46 @@ def _write_epoch_checkpoint(self, svi_state, epoch):
dill.dump(out_dict, f)
os.replace(tmp_file, epoch_file)
+ def restore_svi_from_checkpoint(self, checkpoint_file, init_params=None):
+ """
+ Rebuild a component SVI object and restore its state from a checkpoint,
+ without running any optimization steps or convergence checks.
+
+ ``svi.init()`` must be called at least once to wire up ``constrain_fn``
+ on the SVI object before ``svi.get_params()`` can be used. The state
+ produced by ``init()`` is immediately discarded and replaced with the
+ one loaded from the checkpoint.
+
+ Parameters
+ ----------
+ checkpoint_file : str
+ Path to the checkpoint .pkl file produced by tfs-fit-model.
+ init_params : dict or None, optional
+ Initial parameter values forwarded to ``svi.init()``. The values
+ are never used in inference (the checkpoint overwrites them), but
+ they must be structurally valid for the model.
+
+ Returns
+ -------
+ svi : numpyro.infer.SVI
+ Initialized SVI object with ``constrain_fn`` populated.
+ svi_state : numpyro.infer.svi.SVIState
+ Optimizer state restored from the checkpoint.
+ """
+ svi = self.setup_svi(guide_type="component")
+
+ # init() is required to populate svi.constrain_fn; the resulting state
+ # is thrown away — the checkpoint state is used instead.
+ data_on_gpu = jax.device_put(self.model.data)
+ batch_idx = jax.device_put(self.model.get_random_idx())
+ batch_data = self.model.get_batch(data_on_gpu, batch_idx)
+ init_key = self.get_key()
+ svi.init(init_key, init_params=init_params,
+ priors=self.model.priors, data=batch_data)
+
+ svi_state = self._restore_checkpoint(checkpoint_file)
+ return svi, svi_state
+
def _restore_checkpoint(self,checkpoint_file):
"""
Load an SVI state and PRNG key from a checkpoint file.
diff --git a/src/tfscreen/tfmodel/model_orchestrator.py b/src/tfscreen/tfmodel/model_orchestrator.py
index feed17cc..564705e5 100644
--- a/src/tfscreen/tfmodel/model_orchestrator.py
+++ b/src/tfscreen/tfmodel/model_orchestrator.py
@@ -12,6 +12,7 @@
DataClass,
BindingData,
GrowthData,
+ PreSplitData,
PriorsClass,
GrowthPriors,
BindingPriors
@@ -81,6 +82,12 @@ def _read_growth_df(growth_df,
growth_df = tfscreen.util.io.read_dataframe(growth_df)
growth_df = tfscreen.util.dataframe.get_scaled_cfu(growth_df,need_columns=["ln_cfu","ln_cfu_std"])
+ # Ensure time columns are float so that merge/concat with float prediction
+ # grids (e.g. linspace t_sel) doesn't produce int/float dtype warnings.
+ for _col in ("t_pre", "t_sel"):
+ if _col in growth_df.columns:
+ growth_df[_col] = growth_df[_col].astype(float)
+
# make a replicate column if not defined
if "replicate" not in growth_df.columns:
growth_df["replicate"] = 1
@@ -243,14 +250,24 @@ def _read_binding_df(binding_df,
cols = ["genotype","titrant_name"]
binding_seen = binding_df[cols].drop_duplicates().set_index(cols)
growth_seen = growth_df[cols].drop_duplicates().set_index(cols)
- is_subset = binding_seen.index.isin(growth_seen.index).all()
- if not is_subset:
- raise ValueError(
- "binding_df contains genotype/titrant_name pairs that were not seen "
- "in the growth_df."
+ extra = binding_seen.index[~binding_seen.index.isin(growth_seen.index)]
+ if len(extra) > 0:
+ by_titrant = {}
+ for geno, tname in extra:
+ by_titrant.setdefault(tname, []).append(geno)
+ detail = "\n".join(
+ f" {tname}: {', '.join(sorted(genos))}"
+ for tname, genos in sorted(by_titrant.items())
+ )
+ print(
+ f"Syncing binding_df to growth_df: {len(extra)} "
+ f"genotype/titrant_name pair(s) not in growth_df will be dropped.\n"
+ + detail,
+ flush=True,
)
# Inherit map_theta_group indices from growth_df so parameter indexing
- # stays consistent across the two datasets.
+ # stays consistent across the two datasets. Rows in binding_df with no
+ # match in growth_df are silently dropped by add_group_columns.
binding_df = add_group_columns(binding_df, theta_group_cols,
"map_theta_group", existing_df=growth_df)
else:
@@ -313,6 +330,111 @@ def _build_binding_tm(binding_df):
return binding_tm
+def _read_presplit_df(presplit_df, growth_df):
+ """
+ Read and validate the optional pre-split DataFrame.
+
+ Rows whose genotype does not appear in growth_df are silently dropped.
+ Genotypes present in growth_df but absent from presplit_df are kept; they
+ will produce NaN entries in the tensor (masked out by good_mask).
+
+ Parameters
+ ----------
+ presplit_df : pd.DataFrame or str
+ DataFrame or path to file with pre-split data. Required columns:
+ ``replicate``, ``condition_pre``, ``genotype``, ``ln_cfu``,
+ ``ln_cfu_std``.
+ growth_df : pd.DataFrame
+ The already-processed growth DataFrame (growth_tm.df). Used to
+ determine the valid set of genotypes, replicates, and condition_pre
+ values.
+
+ Returns
+ -------
+ pd.DataFrame
+ A copy of the processed pre-split DataFrame restricted to genotypes
+ present in growth_df.
+ """
+ presplit_df = tfscreen.util.io.read_dataframe(presplit_df)
+ presplit_df = tfscreen.genetics.set_categorical_genotype(presplit_df,
+ standardize=True)
+ tfscreen.util.dataframe.check_columns(
+ presplit_df,
+ required_columns=["replicate", "condition_pre", "genotype",
+ "ln_cfu", "ln_cfu_std"],
+ )
+
+ growth_genotypes = set(growth_df["genotype"])
+ mask = presplit_df["genotype"].isin(growth_genotypes)
+ dropped_genos = sorted(set(presplit_df.loc[~mask, "genotype"].astype(str)))
+ n_dropped = int((~mask).sum())
+ presplit_df = presplit_df[mask].copy()
+ if n_dropped > 0:
+ print(
+ f"Syncing presplit_df to growth_df: {n_dropped} row(s) for "
+ f"{len(dropped_genos)} genotype(s) not found in growth_df will be dropped.\n"
+ f" Dropped genotypes: {', '.join(dropped_genos)}",
+ flush=True,
+ )
+
+ return presplit_df
+
+
+def _build_presplit_tm(presplit_df, growth_tm):
+ """
+ Build a TensorManager for the pre-split (t = -t_pre) data.
+
+ The output tensor has shape ``(num_replicate, num_condition_pre,
+ num_genotype)`` using the *same* categorical orderings as growth_tm so
+ that the genotype axis aligns with ``growth_tm.batch_idx``. Entries
+ with no presplit observation are filled with 1.0 (masked out by
+ good_mask).
+
+ Parameters
+ ----------
+ presplit_df : pd.DataFrame
+ The processed pre-split DataFrame from ``_read_presplit_df``.
+ growth_tm : TensorManager
+ The fully-built growth TensorManager (after ``create_tensors()``).
+
+ Returns
+ -------
+ TensorManager
+ A fully-built TensorManager with tensors ``ln_cfu``, ``ln_cfu_std``,
+ and ``good_mask``.
+ """
+ # Inherit categorical orderings from growth_tm so indices match exactly.
+ dim_names = growth_tm.tensor_dim_names
+ rep_cats = growth_tm.tensor_dim_labels[dim_names.index("replicate")]
+ cp_cats = growth_tm.tensor_dim_labels[dim_names.index("condition_pre")]
+ geno_cats = growth_tm.tensor_dim_labels[dim_names.index("genotype")]
+
+ presplit_df = presplit_df.copy()
+ presplit_df["replicate"] = pd.Categorical(
+ presplit_df["replicate"], categories=rep_cats
+ )
+ presplit_df["condition_pre"] = pd.Categorical(
+ presplit_df["condition_pre"], categories=cp_cats
+ )
+ presplit_df["genotype"] = pd.Categorical(
+ presplit_df["genotype"], categories=geno_cats
+ )
+
+ presplit_tm = TensorManager(presplit_df)
+ presplit_tm.add_pivot_index(tensor_dim_name="replicate",
+ cat_column="replicate")
+ presplit_tm.add_pivot_index(tensor_dim_name="condition_pre",
+ cat_column="condition_pre")
+ presplit_tm.add_pivot_index(tensor_dim_name="genotype",
+ cat_column="genotype")
+
+ presplit_tm.add_data_tensor("ln_cfu", dtype=FLOAT_DTYPE)
+ presplit_tm.add_data_tensor("ln_cfu_std", dtype=FLOAT_DTYPE)
+
+ presplit_tm.create_tensors()
+ return presplit_tm
+
+
def _setup_batching(growth_genotypes,
binding_genotypes,
batch_size):
@@ -358,7 +480,7 @@ def _setup_batching(growth_genotypes,
num_not_binding = len(not_binding_idx)
not_binding_batch_size = batch_size - num_binding
- # Build idx array. The first entries correspond to the binding data and are
+ # Build idx array. The first entries correspond to the binding data and are
# the same for all rounds
idx = np.zeros(batch_size,dtype=int)
idx[:num_binding] = binding_idx
@@ -464,10 +586,12 @@ def __init__(self,
growth_shares_replicates=False,
epistasis=False,
thermo_data=None,
- binding_weight=None):
+ binding_weight=None,
+ presplit_df=None):
self._ln_cfu_df = growth_df
self._binding_df = binding_df
+ self._presplit_df = presplit_df
self._batch_size = batch_size
self._binding_only = binding_only
@@ -805,13 +929,43 @@ def _initialize_data(self):
binding_dataclass = populate_dataclass(BindingData,
sources=binding_data_sources)
+ # ---------------------------------------------------------------------
+ # presplit dataclass (optional)
+
+ if self._presplit_df is not None:
+ self.presplit_df = _read_presplit_df(self._presplit_df,
+ self.growth_tm.df)
+ self.presplit_tm = _build_presplit_tm(self.presplit_df,
+ self.growth_tm)
+ ps_tensors = {
+ "ln_cfu_t0": self.presplit_tm.tensors["ln_cfu"],
+ "ln_cfu_t0_std": self.presplit_tm.tensors["ln_cfu_std"],
+ "good_mask": self.presplit_tm.tensors["good_mask"],
+ }
+ ps_sizes = {
+ "num_replicate": self.presplit_tm.tensor_shape[0],
+ "num_condition_pre": self.presplit_tm.tensor_shape[1],
+ "num_genotype": self.presplit_tm.tensor_shape[2],
+ }
+ presplit_dataclass = populate_dataclass(PreSplitData,
+ sources=[ps_tensors,
+ ps_sizes])
+ else:
+ self.presplit_df = None
+ self.presplit_tm = None
+ presplit_dataclass = None
+
+ # ---------------------------------------------------------------------
+ # Populate DataClass
+
source_data = [{"growth":growth_dataclass,
"binding":binding_dataclass,
+ "presplit":presplit_dataclass,
"num_genotype":self.growth_tm.tensor_shape[-1]}]
source_data.append(full_batch_data)
- # Build the aggregated `DataClass` flax dataclass with the growth
- # and binding dataclasses. Expose as a model attribute.
+ # Build the aggregated `DataClass` flax dataclass with the growth,
+ # binding, and presplit dataclasses. Expose as a model attribute.
self._data = populate_dataclass(DataClass,
sources=source_data)
@@ -1045,11 +1199,20 @@ def _initialize_classes(self):
else:
component_data = self._data.growth
- # Record priors. Components may optionally accept a `data`
- # parameter to derive empirical prior values (e.g. fixed
- # subgroup scales) from the observed data.
+ # Record priors. Components may optionally accept special
+ # keyword arguments:
+ # condition_labels — ordered condition names used by the
+ # linear growth component to set per-condition m priors.
+ # data — component data pytree for empirical prior values.
priors_sig = inspect.signature(component_module.get_priors)
- if "data" in priors_sig.parameters:
+ if "condition_labels" in priors_sig.parameters:
+ cond_rep_df = self.growth_tm.map_groups["condition_rep"]
+ condition_labels = list(
+ cond_rep_df.sort_values("map_condition_rep")["condition_rep"]
+ )
+ priors_class_kwargs[prior_group][key] = \
+ component_module.get_priors(condition_labels=condition_labels)
+ elif "data" in priors_sig.parameters:
priors_class_kwargs[prior_group][key] = \
component_module.get_priors(data=component_data)
else:
@@ -1271,4 +1434,5 @@ def settings(self):
"epistasis": self._epistasis,
"thermo_data": self._thermo_data,
"binding_weight": self._binding_weight,
+ "presplit_df": getattr(self, "_presplit_df", None),
}
diff --git a/src/tfscreen/tfmodel/scripts/configure_model_cli.py b/src/tfscreen/tfmodel/scripts/configure_model_cli.py
index d1df0a33..e52348a7 100644
--- a/src/tfscreen/tfmodel/scripts/configure_model_cli.py
+++ b/src/tfscreen/tfmodel/scripts/configure_model_cli.py
@@ -48,6 +48,7 @@ def check_component_compatibility(condition_growth_model, theta_rescale_model):
def configure_model(binding_df,
growth_df=None,
+ presplit_df=None,
out_prefix="tfs_configure",
condition_growth_model="linear",
growth_transition_model="instant",
@@ -187,8 +188,9 @@ def configure_model(binding_df,
check_component_compatibility(condition_growth_model, theta_rescale_model)
# Initialize model to build mappings and get guesses
- gm = ModelOrchestrator(growth_df,
+ orchestrator = ModelOrchestrator(growth_df,
binding_df,
+ presplit_df=presplit_df,
binding_only=binding_only,
condition_growth=condition_growth_model,
growth_transition=growth_transition_model,
@@ -211,15 +213,18 @@ def configure_model(binding_df,
# Write the model configuration to a file. This includes the model component
# names, the data file paths, and the parameter guesses/priors.
growth_path = None if binding_only else (growth_df if isinstance(growth_df, str) else "growth.csv")
- write_configuration(gm=gm,
+ presplit_path = presplit_df if isinstance(presplit_df, str) else None
+ write_configuration(orchestrator=orchestrator,
out_prefix=out_prefix,
growth_df_path=growth_path,
- binding_df_path=binding_df if isinstance(binding_df, str) else "binding.csv")
+ binding_df_path=binding_df if isinstance(binding_df, str) else "binding.csv",
+ presplit_df_path=presplit_path)
def main():
return generalized_main(configure_model,
manual_arg_types={"binding_df":str,
"growth_df":str,
+ "presplit_df":str,
"spiked":list,
"thermo_data":str,
"batch_size":int,
diff --git a/src/tfscreen/tfmodel/scripts/diagnose_nan_cli.py b/src/tfscreen/tfmodel/scripts/diagnose_nan_cli.py
index 78a0f911..8846da9d 100644
--- a/src/tfscreen/tfmodel/scripts/diagnose_nan_cli.py
+++ b/src/tfscreen/tfmodel/scripts/diagnose_nan_cli.py
@@ -82,8 +82,8 @@ def diagnose_nan(config_file,
jax.config.update("jax_debug_nans", True)
jax.config.update("jax_disable_jit", True)
- gm, init_params = read_configuration(config_file)
- ri = RunInference(gm, seed=seed)
+ orchestrator, init_params = read_configuration(config_file)
+ ri = RunInference(orchestrator, seed=seed)
svi_obj = ri.setup_svi(adam_step_size=1e-6, guide_type="component")
diff --git a/src/tfscreen/tfmodel/scripts/extract_params_cli.py b/src/tfscreen/tfmodel/scripts/extract_params_cli.py
index 970251eb..cc8154df 100644
--- a/src/tfscreen/tfmodel/scripts/extract_params_cli.py
+++ b/src/tfscreen/tfmodel/scripts/extract_params_cli.py
@@ -14,12 +14,13 @@ def extract_params(config_file, param_file, out_prefix="tfs_params"):
If ``param_file`` ends with ``.pkl`` it is treated as a MAP (AutoDelta)
checkpoint: parameters are converted to constrained space and written with
- a single ``point_est`` column.
+ a single ``q0.5`` column.
Otherwise it is treated as a posterior samples ``.h5`` or ``.npz`` file
- produced by tfs-sample-posterior, and the default quantile set (min,
- lower_95, lower_std, lower_quartile, median, upper_quartile, upper_std,
- upper_95, max) is computed.
+ produced by tfs-sample-posterior, and 17 quantile columns are written:
+ ``q0.001``, ``q0.005``, ``q0.01``, ``q0.025``, ``q0.05``, ``q0.1``,
+ ``q0.159``, ``q0.25``, ``q0.5``, ``q0.75``, ``q0.841``, ``q0.9``,
+ ``q0.95``, ``q0.975``, ``q0.99``, ``q0.995``, ``q0.999``.
Writes one CSV per parameter group named ``{out_prefix}_{param_name}.csv``
(e.g. tfs_params_activity.csv, tfs_params_theta.csv).
@@ -34,15 +35,15 @@ def extract_params(config_file, param_file, out_prefix="tfs_params"):
out_prefix : str, optional
Prefix for output CSV files. Default ``'tfs_params'``.
"""
- gm, _ = read_configuration(config_file)
+ orchestrator, _ = read_configuration(config_file)
if param_file.endswith(".pkl"):
- _extract_from_checkpoint(gm, param_file, out_prefix)
+ _extract_from_checkpoint(orchestrator, param_file, out_prefix)
else:
- _extract_from_posteriors(gm, param_file, out_prefix)
+ _extract_from_posteriors(orchestrator, param_file, out_prefix)
-def _extract_from_checkpoint(gm, ckpt_path, out_prefix):
+def _extract_from_checkpoint(orchestrator, ckpt_path, out_prefix):
if not os.path.isfile(ckpt_path):
raise FileNotFoundError(f"Checkpoint file not found: '{ckpt_path}'")
@@ -56,7 +57,7 @@ def _extract_from_checkpoint(gm, ckpt_path, out_prefix):
"tfs-extract-params only supports MAP (AutoDelta) checkpoints."
)
- ri = RunInference(gm, seed=0)
+ ri = RunInference(orchestrator, seed=0)
temp_svi = ri.setup_svi(guide_type="delta")
map_params = temp_svi.optim.get_params(chk_data["svi_state"].optim_state)
@@ -70,7 +71,7 @@ def _extract_from_checkpoint(gm, ckpt_path, out_prefix):
posteriors = {k: np.expand_dims(np.asarray(v), 0) for k, v in constrained.items()}
print(f"Writing parameter CSVs to {out_prefix}_*.csv...", flush=True)
- param_dfs = extract_parameters(gm, posteriors, q_to_get={"point_est": 0.5})
+ param_dfs = extract_parameters(orchestrator, posteriors, q_to_get=[0.5])
for p_name, df in param_dfs.items():
out_file = f"{out_prefix}_{p_name}.csv"
df.to_csv(out_file, index=False)
@@ -78,7 +79,7 @@ def _extract_from_checkpoint(gm, ckpt_path, out_prefix):
print("Done.", flush=True)
-def _extract_from_posteriors(gm, posterior_file, out_prefix):
+def _extract_from_posteriors(orchestrator, posterior_file, out_prefix):
if not os.path.exists(posterior_file):
if os.path.exists(f"{posterior_file}.h5"):
posterior_file = f"{posterior_file}.h5"
@@ -103,7 +104,7 @@ def _extract_from_posteriors(gm, posterior_file, out_prefix):
"Extraction of natural parameters may fail.", flush=True)
print(f"Extracting parameters to {out_prefix}_*.csv...", flush=True)
- params = extract_parameters(gm, posteriors)
+ params = extract_parameters(orchestrator, posteriors)
for p_name, p_df in params.items():
p_df.to_csv(f"{out_prefix}_{p_name}.csv", index=False)
diff --git a/src/tfscreen/tfmodel/scripts/fit_model_cli.py b/src/tfscreen/tfmodel/scripts/fit_model_cli.py
index 0b78028f..15c505bf 100644
--- a/src/tfscreen/tfmodel/scripts/fit_model_cli.py
+++ b/src/tfscreen/tfmodel/scripts/fit_model_cli.py
@@ -24,10 +24,6 @@ def _run_map(ri,
convergence_check_interval=2,
checkpoint_interval=10,
max_num_epochs=100000,
- num_posterior_samples=10000,
- sampling_batch_size=100,
- forward_batch_size=512,
- always_get_posterior=False,
init_param_jitter=0.0,
epoch_checkpoint_interval=1000):
"""
@@ -76,6 +72,11 @@ def _run_map(ri,
``checkpoints/`` subdirectory (default 1000). Set to 0 or None to
disable.
+ Notes
+ -----
+ Posterior sampling is not performed here. Call ``tfs-sample-posterior``
+ after fitting to draw posterior samples.
+
Returns
-------
svi_state : Any
@@ -118,15 +119,6 @@ def _run_map(ri,
# Write the current parameter values
ri.write_params(params,out_prefix=out_prefix)
- if always_get_posterior:
-
- ri.get_posteriors(svi=map_obj,
- svi_state=svi_state,
- out_prefix=out_prefix,
- num_posterior_samples=num_posterior_samples,
- sampling_batch_size=sampling_batch_size,
- forward_batch_size=forward_batch_size)
-
# Write convergence information to stdout
if converged:
print("MAP run converged.",flush=True)
@@ -149,10 +141,6 @@ def _run_svi(ri,
convergence_check_interval=2,
checkpoint_interval=10,
max_num_epochs=100000,
- num_posterior_samples=10000,
- sampling_batch_size=100,
- forward_batch_size=512,
- always_get_posterior=False,
init_param_jitter=0.1,
epoch_checkpoint_interval=1000):
"""
@@ -194,16 +182,6 @@ def _run_svi(ri,
Frequency (in epochs) between checkpoints (default 10).
max_num_epochs : int, optional
Maximum number of SVI optimization epochs (default 100000).
- num_posterior_samples : int, optional
- Number of posterior samples to draw after convergence (default 10000).
- sampling_batch_size : int, optional
- When getting posteriors, sample parameter posteriors in batches of this
- size (default 100).
- forward_batch_size : int, optional
- when getting posteriors, calculate forward predictions in batches of
- this size (default 512)
- always_get_posterior : bool, optional
- If True, always sample posteriors even if not converged (default False).
init_param_jitter : float, optional
Amount of jitter to add to init_params (default 0.1).
epoch_checkpoint_interval : int or None, optional
@@ -211,8 +189,17 @@ def _run_svi(ri,
``checkpoints/`` subdirectory (default 1000). Set to 0 or None to
disable.
+ Notes
+ -----
+ Posterior sampling is not performed here. Call ``tfs-sample-posterior``
+ after fitting to draw posterior samples.
+
Returns
-------
+ svi_obj : Any
+ The SVI optimizer object created by ``ri.setup_svi``. Returned so
+ that callers (e.g. ``tfs-sample-posterior``) can pass it directly to
+ ``ri.get_posteriors`` without needing to re-create the guide.
svi_state : Any
Final optimizer state object from SVI.
svi_params : dict
@@ -251,22 +238,15 @@ def _run_svi(ri,
epoch_checkpoint_interval=epoch_checkpoint_interval
)
- if converged or always_get_posterior:
-
- ri.get_posteriors(svi=svi_obj,
- svi_state=svi_state,
- out_prefix=out_prefix,
- num_posterior_samples=num_posterior_samples,
- sampling_batch_size=sampling_batch_size,
- forward_batch_size=forward_batch_size)
+ # Write convergence information to stdout (skip when restoring from
+ # checkpoint with no additional epochs — convergence is not meaningful).
+ if max_num_epochs > 0:
+ if converged:
+ print("SVI run converged.",flush=True)
+ else:
+ print("SVI run has not yet converged.",flush=True)
- # Write convergence information to stdout
- if converged:
- print("SVI run converged.",flush=True)
- else:
- print("SVI run has not yet converged.",flush=True)
-
- return svi_state, params, converged
+ return svi_obj, svi_state, params, converged
def _run_nuts(ri,
out_prefix="tfs",
@@ -340,10 +320,7 @@ def fit_model(config_file,
convergence_check_interval=2,
checkpoint_interval=10,
max_num_epochs=100000,
- num_posterior_samples=10000,
- sampling_batch_size=100,
forward_batch_size=512,
- always_get_posterior=False,
pre_map_num_epoch=1000,
init_param_jitter=0.1,
nuts_num_warmup=500,
@@ -368,7 +345,8 @@ def fit_model(config_file,
Path to a checkpoint file to resume SVI from, or None to start fresh.
analysis_method : str, optional
Method for inference. Allowed values are 'svi' (default), 'map', or 'nuts'.
- To draw posterior samples from an existing checkpoint, use tfs-sample-posterior.
+ Case-insensitive. Posterior sampling is not performed; call
+ ``tfs-sample-posterior`` after fitting.
out_prefix : str, optional
Prefix for all output files: checkpoints, parameter files, and the
posterior HDF5 (default 'tfs_fit_model'). Files are named
@@ -394,17 +372,9 @@ def fit_model(config_file,
Frequency (in epochs) between checkpoints (default 10).
max_num_epochs : int, optional
Maximum number of SVI optimization epochs (default 100000).
- num_posterior_samples : int, optional
- Number of posterior samples to draw after convergence (default 10000).
- sampling_batch_size : int, optional
- When getting posteriors, sample parameter posteriors in batches of this
- size (default 100).
forward_batch_size : int, optional
- When getting posteriors, calculate forward predictions in batches of
- this size (default 512).
- always_get_posterior : bool, optional
- If True, always generate and save posterior samples, even if the
- optimization did not formally converge (default False).
+ When getting NUTS posteriors, calculate forward predictions in batches
+ of this size (default 512).
pre_map_num_epoch : int, optional
Number of MAP iterations to run prior to SVI (default 1000). Only used
if analysis_method is 'svi'.
@@ -442,6 +412,8 @@ def fit_model(config_file,
if seed is None and checkpoint_file is None:
raise ValueError("seed must be provided unless loading from a checkpoint.")
+ analysis_method = analysis_method.lower()
+
# Check for existing results to avoid overwriting unless resuming
if checkpoint_file is None:
checkpoint_path = f"{out_prefix}_checkpoint.pkl"
@@ -460,14 +432,14 @@ def fit_model(config_file,
"overwrite, delete the file or change out_prefix."
)
- gm, init_params = read_configuration(config_file)
+ orchestrator, init_params = read_configuration(config_file)
# For posterior mode the seed is optional: the checkpoint restores the PRNG
# key for SVI checkpoints, and any valid key works for MAP/Laplace sampling.
effective_seed = seed if seed is not None else 0
# Run SVI / MAP
- ri = RunInference(gm, effective_seed)
+ ri = RunInference(orchestrator, effective_seed)
if analysis_method == "svi":
if pre_map_num_epoch > 0 and checkpoint_file is None:
@@ -498,10 +470,6 @@ def fit_model(config_file,
convergence_check_interval=convergence_check_interval,
checkpoint_interval=checkpoint_interval,
max_num_epochs=max_num_epochs,
- num_posterior_samples=num_posterior_samples,
- sampling_batch_size=sampling_batch_size,
- forward_batch_size=forward_batch_size,
- always_get_posterior=always_get_posterior,
init_param_jitter=init_param_jitter,
epoch_checkpoint_interval=epoch_checkpoint_interval)
@@ -520,10 +488,6 @@ def fit_model(config_file,
convergence_check_interval=convergence_check_interval,
checkpoint_interval=checkpoint_interval,
max_num_epochs=max_num_epochs,
- num_posterior_samples=num_posterior_samples,
- sampling_batch_size=sampling_batch_size,
- forward_batch_size=forward_batch_size,
- always_get_posterior=always_get_posterior,
init_param_jitter=init_param_jitter,
epoch_checkpoint_interval=epoch_checkpoint_interval)
diff --git a/src/tfscreen/tfmodel/scripts/predict_growth_cli.py b/src/tfscreen/tfmodel/scripts/predict_growth_cli.py
index 389fcd8c..72d4f0ed 100644
--- a/src/tfscreen/tfmodel/scripts/predict_growth_cli.py
+++ b/src/tfscreen/tfmodel/scripts/predict_growth_cli.py
@@ -6,7 +6,7 @@
def predict_growth(config_file,
param_file,
- out_prefix="tfs_growth_pred",
+ out_prefix="tfs_pred_growth",
genotypes_file=None,
titrant_names_file=None,
titrant_concs_file=None,
@@ -51,7 +51,7 @@ def predict_growth(config_file,
tfs-sample-posterior first.
out_prefix : str, optional
Prefix for the output CSV file. Written to {out_prefix}.csv.
- Default 'tfs_growth_pred'.
+ Default 'tfs_pred_growth'.
genotypes_file : str or None, optional
Plain-text file with one genotype per line (slash-separated mutations,
e.g. 'M42I/K84L', or 'wt'). These genotypes are unioned with all
@@ -80,29 +80,29 @@ def predict_growth(config_file,
file_concs = [float(x) for x in read_lines(titrant_concs_file)] if titrant_concs_file else []
print(f"Loading configuration from {config_file}...", flush=True)
- gm, _ = read_configuration(config_file)
+ orchestrator, _ = read_configuration(config_file)
is_map = param_file.endswith(".pkl")
- param_file = resolve_param_file(param_file, gm, out_prefix)
+ param_file = resolve_param_file(param_file, orchestrator, out_prefix)
# Build training-data membership set for in_training_data column.
training_tuples = set(
- zip(gm.growth_df["genotype"],
- gm.growth_df["titrant_name"],
- gm.growth_df["titrant_conc"])
+ zip(orchestrator.growth_df["genotype"],
+ orchestrator.growth_df["titrant_name"],
+ orchestrator.growth_df["titrant_conc"])
)
if only_files:
genotypes = file_genotypes if file_genotypes else None
titrant_concs = file_concs if file_concs else None
else:
- training_genotypes = list(gm.growth_df["genotype"].unique())
+ training_genotypes = list(orchestrator.growth_df["genotype"].unique())
genotypes = list(dict.fromkeys(training_genotypes + file_genotypes)) if file_genotypes else None
- training_concs = list(gm.growth_df["titrant_conc"].unique())
+ training_concs = list(orchestrator.growth_df["titrant_conc"].unique())
titrant_concs = sorted(set(training_concs) | set(file_concs)) if file_concs else None
- q_to_get = {"point_est": 0.5} if is_map else None
+ q_to_get = [0.5] if is_map else None
print("Running growth predictions...", flush=True)
- result_df = predict(model_class=gm,
+ result_df = predict(orchestrator=orchestrator,
param_posteriors=param_file,
predict_sites=["growth_pred"],
num_samples=num_samples,
diff --git a/src/tfscreen/tfmodel/scripts/predict_theta_cli.py b/src/tfscreen/tfmodel/scripts/predict_theta_cli.py
index 3f89069e..5e7b04a2 100644
--- a/src/tfscreen/tfmodel/scripts/predict_theta_cli.py
+++ b/src/tfscreen/tfmodel/scripts/predict_theta_cli.py
@@ -11,7 +11,7 @@
def predict_theta(config_file,
param_file,
- out_prefix="tfs_theta_pred",
+ out_prefix="tfs_pred_theta",
genotypes_file=None,
titrant_names_file=None,
titrant_concs_file=None,
@@ -60,7 +60,7 @@ def predict_theta(config_file,
tfs-sample-posterior first.
out_prefix : str, optional
Prefix for the output CSV file. Written to {out_prefix}.csv.
- Default 'tfs_theta_pred'.
+ Default 'tfs_pred_theta'.
genotypes_file : str or None, optional
Plain-text file with one genotype per line (slash-separated mutations,
e.g. 'M42I/K84L', or 'wt'). These genotypes are unioned with all
@@ -97,18 +97,18 @@ def predict_theta(config_file,
)
print(f"Loading configuration from {config_file}...", flush=True)
- gm, _ = read_configuration(config_file)
+ orchestrator, _ = read_configuration(config_file)
is_map = param_file.endswith(".pkl")
- param_file = resolve_param_file(param_file, gm, out_prefix)
+ param_file = resolve_param_file(param_file, orchestrator, out_prefix)
# Determine training genotypes and (genotype, titrant_name, titrant_conc) set.
# growth_tm is preferred (more genotypes); binding_tm is the fallback for
# binding-only runs where growth_tm is None.
- training_genotypes = set(gm.training_tm.df["genotype"].unique())
+ training_genotypes = set(orchestrator.training_tm.df["genotype"].unique())
training_tuples = set(
- zip(gm.training_tm.df["genotype"],
- gm.training_tm.df["titrant_name"],
- gm.training_tm.df["titrant_conc"])
+ zip(orchestrator.training_tm.df["genotype"],
+ orchestrator.training_tm.df["titrant_name"],
+ orchestrator.training_tm.df["titrant_conc"])
)
# Resolve requested genotypes.
@@ -141,7 +141,7 @@ def predict_theta(config_file,
manual_titrant_df = file_titrant_df
else:
training_titrant_df = (
- gm.training_tm.df[["titrant_name", "titrant_conc"]]
+ orchestrator.training_tm.df[["titrant_name", "titrant_conc"]]
.drop_duplicates()
.reset_index(drop=True)
)
@@ -157,10 +157,10 @@ def predict_theta(config_file,
out_of_training = [g for g in requested_genotypes if g not in training_genotypes]
if out_of_training:
- module = model_registry.get("theta", {}).get(gm._theta)
+ module = model_registry.get("theta", {}).get(orchestrator._theta)
if module is None or not hasattr(module, "predict_unmeasured"):
raise ValueError(
- f"The theta component '{gm._theta}' does not support prediction "
+ f"The theta component '{orchestrator._theta}' does not support prediction "
"for genotypes not seen during training. Remove out-of-training "
"genotypes from genotypes_file, or use a theta component that "
"implements predict_unmeasured (e.g. 'hill')."
@@ -170,13 +170,13 @@ def predict_theta(config_file,
if manual_titrant_df is None:
# Use unique (titrant_name, titrant_conc) pairs from training data.
manual_titrant_df = (
- gm.training_tm.df[["titrant_name", "titrant_conc"]]
+ orchestrator.training_tm.df[["titrant_name", "titrant_conc"]]
.drop_duplicates()
.reset_index(drop=True)
)
- q_to_get = {"point_est": 0.5} if is_map else None
+ q_to_get = [0.5] if is_map else None
result_df = extract_theta_unmeasured(
- model=gm,
+ orchestrator=orchestrator,
posteriors=param_file,
target_genotypes=requested_genotypes,
manual_titrant_df=manual_titrant_df,
@@ -186,9 +186,9 @@ def predict_theta(config_file,
else:
print(f"Predicting theta for {len(requested_genotypes)} training genotype(s)...",
flush=True)
- q_to_get = {"point_est": 0.5} if is_map else None
+ q_to_get = [0.5] if is_map else None
result_df = extract_theta_curves(
- model=gm,
+ orchestrator=orchestrator,
posteriors=param_file,
manual_titrant_df=manual_titrant_df,
num_samples=num_samples,
diff --git a/src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py b/src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py
similarity index 51%
rename from src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py
rename to src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py
index fbc1a090..7c0883c9 100644
--- a/src/tfscreen/tfmodel/scripts/run_prefit_calibration_cli.py
+++ b/src/tfscreen/tfmodel/scripts/prefit_calibration_cli.py
@@ -11,7 +11,9 @@
- ``theta`` → ``simple`` (theta values pinned from binding data)
- ``activity`` → ``hierarchical`` with hyperparams *pinned* to their
prior locs (degenerate, no learning)
-- ``dk_geno`` → ``hierarchical`` with hyperparams *pinned*
+- ``dk_geno`` → ``fixed`` (all zeros; eliminates the WT-anchor
+ bias that arises when mutants have non-zero mean
+ pleiotropic effects)
- ``ln_cfu0`` → ``hierarchical`` with hyperparams *pinned*
- ``transformation`` → ``single`` (no learning)
- ``theta_*_noise`` → ``zero`` (no learning)
@@ -22,23 +24,17 @@
After MAP convergence, the script:
-1. Estimates a per-site 1-sigma uncertainty from the Hessian of the
- negative log-joint at the MAP point (delta-method propagated through
- any constrained-support bijections).
-2. Identifies which fields of the production ``ModelPriors`` provide
+1. Identifies which fields of the production ``ModelPriors`` provide
``dist.loc`` / ``dist.scale`` for each calibrated sample site, using
a sentinel-trace introspection so we don't have to hard-code the
field-naming conventions of every component.
-3. Updates the production ``{out_prefix}_guesses.csv`` *in place*, writing
- a ``.bak`` backup before overwriting. For simple-prior
+2. Updates the production priors and guesses CSVs *in place*, writing
+ ``.bak`` backups before overwriting. For simple-prior
``condition_growth`` and ``growth_transition`` components, per-condition
MAP estimates are written directly as per-condition guess values
(``{site}_locs`` rows in the guesses CSV), giving the production SVI a
warm start from the calibration fit. Only rows belonging to
``condition_growth`` or ``growth_transition`` are touched.
-
- For legacy hierarchical components (if any), scalar hyper-site MAP
- estimates also update the corresponding rows in the priors CSV.
"""
import dataclasses
@@ -75,7 +71,7 @@
_CALIBRATION_OVERRIDES = {
"theta": "_simple",
"activity": "hierarchical_geno",
- "dk_geno": "hierarchical_geno",
+ "dk_geno": "fixed",
"ln_cfu0": "hierarchical",
"transformation": "single",
"theta_growth_noise": "zero",
@@ -92,14 +88,9 @@
("hyper_loc", "hyper_loc_loc"),
("hyper_scale", "hyper_scale_loc"),
),
- "dk_geno": (
- ("hyper_loc", "hyper_loc_loc"),
- ("hyper_scale", "hyper_scale_loc"),
- ("hyper_shift", "hyper_shift_loc"),
- ),
"ln_cfu0": (
- ("hyper_loc", "ln_cfu0_hyper_loc_loc"),
- ("hyper_scale", "ln_cfu0_hyper_scale_loc"),
+ ("hyper_loc", "ln_cfu0_hyper_loc_locs"),
+ ("hyper_scale", "ln_cfu0_hyper_scale_locs"),
),
}
@@ -107,6 +98,20 @@
# downstream logit transform (in the simple-theta component) is finite.
_THETA_EPS = 1e-6
+# Default floors for Hessian-derived prior scales. These represent the
+# minimum plausible inter-experiment variability in growth-rate parameters
+# that the calibration MAP cannot see (e.g., inoculum variation, day-to-day
+# media differences). For real data, calibrate these by running the prefit
+# on multiple independent replicates and examining the spread of MAP estimates.
+_DEFAULT_K_SCALE_FLOOR = 0.002 # hr⁻¹ — baseline growth rate floor
+_DEFAULT_M_SCALE_FLOOR = 0.001 # hr⁻¹ — theta-slope floor
+# Ceilings match the linear-component default prior scales so the prefit can
+# only TIGHTEN (never loosen) the production prior. A degenerate Hessian
+# direction (e.g. b ↔ ln_cfu0 with fixed t_pre) can give sigma >> 10, which
+# without a ceiling would silently widen the prior far beyond the default.
+_DEFAULT_K_SCALE_CEILING = 0.1 # matches linear.get_hyperparameters() k_scale
+_DEFAULT_M_SCALE_CEILING = 0.01 # matches linear.get_hyperparameters() m_scale_plus
+
# ---------------------------------------------------------------------------
# Data filtering
@@ -173,39 +178,48 @@ def _intersect_data(growth_df, binding_df):
# theta_values for the simple component
# ---------------------------------------------------------------------------
-def _compute_theta_values(gm_cal, binding_df_cal):
+def _compute_theta_values(orchestrator_cal, binding_df_cal):
"""
- Build a (num_titrant_name, num_titrant_conc) theta tensor for the
- simple-theta component of the calibration model.
+ Build a per-genotype ``(T, C, G)`` theta tensor for the simple-theta
+ component of the calibration model.
- Each cell is the inverse-variance weighted mean of ``theta_obs``
- across the genotypes that contributed observations at that
- (titrant_name, titrant_conc) cell. Cells with no usable observations
- fall back to a plain mean of ``theta_obs``; if that is also empty the
- cell is set to 0.5 (uninformative midpoint).
+ Each cell ``[i, j, k]`` holds the inverse-variance weighted mean of
+ ``theta_obs`` for titrant_name ``i``, titrant_conc ``j``, and genotype
+ ``k``. Cells with no usable observations fall back to a plain mean;
+ if that is also empty the cell is set to 0.5 (uninformative midpoint).
The dimension ordering is taken from
- ``gm_cal.binding_tm.tensor_dim_labels`` so the resulting array
- matches the layout the simple-theta component expects.
+ ``orchestrator_cal.binding_tm.tensor_dim_labels`` so the resulting array
+ matches the layout the simple-theta component expects. The genotype
+ axis ordering mirrors the calibration model's internal genotype index so
+ that ``theta_values[:, :, k]`` is the correct curve for the k-th
+ calibration genotype.
"""
- tn_idx = gm_cal.binding_tm.tensor_dim_names.index("titrant_name")
- tc_idx = gm_cal.binding_tm.tensor_dim_names.index("titrant_conc")
- titrant_name_labels = list(gm_cal.binding_tm.tensor_dim_labels[tn_idx])
- titrant_conc_labels = list(gm_cal.binding_tm.tensor_dim_labels[tc_idx])
+ tn_idx = orchestrator_cal.binding_tm.tensor_dim_names.index("titrant_name")
+ tc_idx = orchestrator_cal.binding_tm.tensor_dim_names.index("titrant_conc")
+ geno_idx = orchestrator_cal.binding_tm.tensor_dim_names.index("genotype")
+
+ titrant_name_labels = list(orchestrator_cal.binding_tm.tensor_dim_labels[tn_idx])
+ titrant_conc_labels = list(orchestrator_cal.binding_tm.tensor_dim_labels[tc_idx])
+ genotype_labels = list(orchestrator_cal.binding_tm.tensor_dim_labels[geno_idx])
n_name = len(titrant_name_labels)
n_conc = len(titrant_conc_labels)
- theta_values = np.full((n_name, n_conc), 0.5, dtype=float)
+ n_geno = len(genotype_labels)
+ theta_values = np.full((n_name, n_conc, n_geno), 0.5, dtype=float)
name_to_i = {str(n): i for i, n in enumerate(titrant_name_labels)}
conc_to_j = {float(c): j for j, c in enumerate(titrant_conc_labels)}
+ geno_to_k = {str(g): k for k, g in enumerate(genotype_labels)}
- grouped = binding_df_cal.groupby(["titrant_name", "titrant_conc"],
- observed=True)
- for (tn, tc), grp in grouped:
+ grouped = binding_df_cal.groupby(
+ ["titrant_name", "titrant_conc", "genotype"], observed=True
+ )
+ for (tn, tc, geno), grp in grouped:
i = name_to_i.get(str(tn))
j = conc_to_j.get(float(tc))
- if i is None or j is None:
+ k = geno_to_k.get(str(geno))
+ if i is None or j is None or k is None:
continue
theta_obs = grp["theta_obs"].to_numpy(dtype=float)
@@ -214,11 +228,11 @@ def _compute_theta_values(gm_cal, binding_df_cal):
& (theta_std > 0))
if valid.any():
w = 1.0 / np.square(theta_std[valid])
- theta_values[i, j] = float(np.sum(theta_obs[valid] * w) / np.sum(w))
+ theta_values[i, j, k] = float(np.sum(theta_obs[valid] * w) / np.sum(w))
else:
usable = np.isfinite(theta_obs)
if usable.any():
- theta_values[i, j] = float(np.mean(theta_obs[usable]))
+ theta_values[i, j, k] = float(np.mean(theta_obs[usable]))
np.clip(theta_values, _THETA_EPS, 1.0 - _THETA_EPS, out=theta_values)
return jnp.asarray(theta_values, dtype=jnp.float32)
@@ -228,7 +242,7 @@ def _compute_theta_values(gm_cal, binding_df_cal):
# Calibration model construction
# ---------------------------------------------------------------------------
-def _build_calibration_model(gm_prod, growth_df_cal, binding_df_cal):
+def _build_calibration_model(orchestrator_prod, growth_df_cal, binding_df_cal):
"""
Construct the calibration ``ModelOrchestrator`` with hardcoded calibration
overrides applied on top of the production component selections.
@@ -239,7 +253,7 @@ def _build_calibration_model(gm_prod, growth_df_cal, binding_df_cal):
sees the (calibration-genotype × calibration-condition) intersection,
which by construction does not include the production's spiked rows.
"""
- settings = dict(gm_prod.settings)
+ settings = dict(orchestrator_prod.settings)
for k, v in _CALIBRATION_OVERRIDES.items():
settings[k] = v
settings["spiked_genotypes"] = None
@@ -256,7 +270,7 @@ def _build_calibration_model(gm_prod, growth_df_cal, binding_df_cal):
**settings)
-def _inject_calibration_priors(gm_cal, gm_prod, theta_values):
+def _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values):
"""
Wire the calibration model's priors into shape:
@@ -270,10 +284,10 @@ def _inject_calibration_priors(gm_cal, gm_prod, theta_values):
those hyperparameters to their prior loc / scale defaults so the
MAP doesn't try to learn them from calibration-only data.
- Mutates ``gm_cal._priors`` in place.
+ Mutates ``orchestrator_cal._priors`` in place.
"""
- prod_growth = gm_prod.priors.growth
- cal_growth = gm_cal.priors.growth
+ prod_growth = orchestrator_prod.priors.growth
+ cal_growth = orchestrator_cal.priors.growth
cg_prod = prod_growth.condition_growth
cg_cal = cal_growth.condition_growth
@@ -327,7 +341,16 @@ def _inject_calibration_priors(gm_cal, gm_prod, theta_values):
pinned_dict = {}
for suffix, field_name in suffix_field_pairs:
if hasattr(comp, field_name):
- pinned_dict[suffix] = float(getattr(comp, field_name))
+ val = getattr(comp, field_name)
+ arr = np.asarray(val)
+ if arr.ndim == 0:
+ # Scalar field (e.g. activity hyper_loc_loc)
+ pinned_dict[suffix] = float(arr)
+ else:
+ # Array field (e.g. ln_cfu0 per-class arrays) — expand to
+ # per-class entries so _hyper() finds "hyper_loc_0" etc.
+ for i, v in enumerate(arr.flat):
+ pinned_dict[f"{suffix}_{i}"] = float(v)
if pinned_dict and hasattr(comp, "replace"):
pinned_components[comp_name] = comp.replace(pinned=pinned_dict)
@@ -337,21 +360,21 @@ def _inject_calibration_priors(gm_cal, gm_prod, theta_values):
new_growth = cal_growth.replace(**growth_updates)
# theta priors live on PriorsClass.theta
- theta_priors = gm_cal.priors.theta
+ theta_priors = orchestrator_cal.priors.theta
if hasattr(theta_priors, "replace"):
new_theta = theta_priors.replace(theta_values=theta_values)
else:
new_theta = theta_priors
- new_priors = gm_cal.priors.replace(growth=new_growth, theta=new_theta)
- gm_cal._priors = new_priors
+ new_priors = orchestrator_cal.priors.replace(growth=new_growth, theta=new_theta)
+ orchestrator_cal._priors = new_priors
# ---------------------------------------------------------------------------
# Field mapping introspection
# ---------------------------------------------------------------------------
-def _identify_field_mapping(gm_cal):
+def _identify_field_mapping(orchestrator_cal):
"""
Enumerate the ``condition_growth`` and ``growth_transition`` sample sites
and derive their ``ModelPriors`` field names.
@@ -376,8 +399,8 @@ def _identify_field_mapping(gm_cal):
dict[str, dict]
Site name → ``{"component", "dist_class", "is_array", ...field names...}``.
"""
- model_trace = trace(seed(gm_cal.jax_model, rng_seed=0)).get_trace(
- data=gm_cal.data, priors=gm_cal.priors
+ model_trace = trace(seed(orchestrator_cal.jax_model, rng_seed=0)).get_trace(
+ data=orchestrator_cal.data, priors=orchestrator_cal.priors
)
out = {}
@@ -438,31 +461,30 @@ def _csv_row_name(component, field_name):
return f"growth.{component}.{field_name}"
-def _build_csv_updates(field_mapping, hessian_results):
+def _build_csv_updates(field_mapping, params):
"""
- Translate the sentinel-trace mapping plus per-site MAP / sigma
- estimates into two flat dicts of in-place updates:
+ Translate the sentinel-trace field mapping and MAP params into two flat
+ dicts of in-place updates.
Returns
-------
prior_updates : dict[str, float]
``{csv_row_name → new_value}`` for the priors CSV.
- Normal sites contribute two rows (loc field ← MAP, scale field ←
- Hessian sigma); HalfNormal contributes one (scale field ← MAP,
- recentering the prior on the MAP point).
- guess_updates : dict[str, float]
- ``{site_name → MAP value}`` for the guesses CSV. Only scalar
- sites are included.
+ Scalar Normal/HalfNormal sites write the MAP value to their
+ ``loc_field`` / ``scale_field`` respectively.
+ guess_updates : dict[str, np.ndarray | float]
+ ``{site_name → MAP value}`` for the guesses CSV.
+ Array sites write per-condition MAP arrays (keyed as
+ ``{site_name}_locs``); scalar sites write a single float.
"""
prior_updates = {}
guess_updates = {}
for site_name, info in field_mapping.items():
- if site_name not in hessian_results:
+ auto_loc_key = f"{site_name}_auto_loc"
+ if auto_loc_key not in params:
continue
- result = hessian_results[site_name]
- map_val_arr = np.asarray(result["map"])
- sigma_arr = np.asarray(result["sigma"])
+ map_val_arr = np.asarray(params[auto_loc_key])
component = info["component"]
dist_class = info["dist_class"]
@@ -471,9 +493,8 @@ def _build_csv_updates(field_mapping, hessian_results):
is_array = info.get("is_array", False)
if is_array:
- # Simple-prior per-condition array site.
- # Write per-condition MAP estimates to guesses so the production
- # SVI starts from the calibration fit. Priors are left unchanged.
+ # Simple-prior per-condition array site: write MAP array to guesses
+ # so the production SVI starts from the calibration fit.
if loc_field is not None:
guess_updates[f"{site_name}_locs"] = map_val_arr
else:
@@ -481,28 +502,126 @@ def _build_csv_updates(field_mapping, hessian_results):
if map_val_arr.shape != ():
continue
map_val = float(map_val_arr)
- sigma_val = float(sigma_arr) if sigma_arr.shape == () else None
if dist_class == "Normal":
if loc_field is not None:
prior_updates[_csv_row_name(component, loc_field)] = map_val
- if scale_field is not None and sigma_val is not None:
- prior_updates[_csv_row_name(component, scale_field)] = sigma_val
elif dist_class == "HalfNormal":
if scale_field is not None:
- # Recenter the HalfNormal on the MAP point.
prior_updates[_csv_row_name(component, scale_field)] = map_val
else:
if loc_field is not None:
prior_updates[_csv_row_name(component, loc_field)] = map_val
- if scale_field is not None and sigma_val is not None:
- prior_updates[_csv_row_name(component, scale_field)] = sigma_val
guess_updates[site_name] = map_val
return prior_updates, guess_updates
+def _build_hessian_scale_updates(field_mapping, hessian_results,
+ k_scale_floor, m_scale_floor,
+ k_scale_ceiling, m_scale_ceiling):
+ """
+ Build per-condition scale updates from Hessian-derived sigmas with floors
+ and ceilings.
+
+ For each array condition_growth site (``k`` and ``m``) that appears in
+ both ``field_mapping`` and ``hessian_results``, the per-element Hessian
+ sigma is clipped to ``[floor, ceiling]`` and returned in two dicts:
+
+ * **guess_updates** — per-condition sigma arrays keyed as
+ ``{site_name}_scales``. Written to the guesses CSV to initialise the
+ variational posterior's uncertainty near the MAP calibration.
+
+ * **prior_updates** — scalar (max across conditions) written to the priors
+ CSV. For ``k``, updates ``condition_growth.k_scale``; for ``m``,
+ updates ``condition_growth.m_scale_plus`` (selection conditions). This
+ is the constraint that actually prevents the production SVI from drifting
+ away from the calibration.
+
+ **Why both floor and ceiling?**
+
+ The ceiling is essential. The prefit calibration uses ``growth_noise:
+ "zero"`` and a fixed ``t_pre``, which creates a near-degenerate Hessian
+ direction (the baseline growth rate ``b`` and initial count ``ln_cfu0`` are
+ only weakly separated). The clamping at 1e-3 inside
+ ``compute_hessian_sigmas`` then gives ``sigma ≈ 1/sqrt(1e-3) ≈ 31`` for
+ that direction — a huge value that would *loosen* the production prior far
+ beyond the original default. The ceiling prevents this: if the Hessian
+ sigma exceeds the original production prior scale, the original scale is
+ kept.
+
+ The floor prevents over-constraining from the opposite direction: if the
+ calibration data is very informative and the Hessian gives
+ ``sigma ≈ 1e-5``, we still keep at least ``floor`` of freedom to represent
+ irreducible inter-experiment variability.
+
+ Parameters
+ ----------
+ field_mapping : dict
+ Output of :func:`_identify_field_mapping`.
+ hessian_results : dict
+ Output of ``RunInference.compute_hessian_sigmas``; maps site names to
+ ``{"map": array, "sigma": array}`` dicts.
+ k_scale_floor : float
+ Minimum prior scale for baseline growth-rate (k / b) parameters.
+ m_scale_floor : float
+ Minimum prior scale for theta-slope (m) parameters.
+ k_scale_ceiling : float
+ Maximum prior scale for k. Hessian sigmas above this are capped so
+ the prefit never *loosens* the production k prior beyond its default.
+ m_scale_ceiling : float
+ Maximum prior scale for m_scale_plus. Analogous to ``k_scale_ceiling``.
+
+ Returns
+ -------
+ guess_updates : dict[str, np.ndarray]
+ Per-condition clipped sigma arrays for the guesses CSV.
+ prior_updates : dict[str, float]
+ Scalar (max-across-conditions) prior scale updates for the priors CSV.
+ """
+ guess_updates = {}
+ prior_updates = {}
+
+ for site_name, info in field_mapping.items():
+ if not info.get("is_array", False):
+ continue
+ if site_name not in hessian_results:
+ continue
+
+ sigma = np.asarray(hessian_results[site_name]["sigma"])
+ loc_field = info.get("loc_field", "")
+ scale_field = info.get("scale_field")
+ component = info["component"]
+
+ if loc_field == "k_loc":
+ floor = k_scale_floor
+ ceiling = k_scale_ceiling
+ prior_field = scale_field # "k_scale"
+ elif loc_field == "m_loc":
+ floor = m_scale_floor
+ ceiling = m_scale_ceiling
+ prior_field = "m_scale_plus" # tighten the + condition prior
+ else:
+ continue
+
+ # Clip: floor prevents over-constraining; ceiling prevents loosening.
+ clipped_sigma = np.clip(sigma, floor, ceiling)
+
+ # Per-condition scales to guesses CSV (initialise variational uncertainty)
+ if scale_field is not None:
+ guess_updates[f"{site_name}_scales"] = clipped_sigma
+
+ # Scalar (max) to priors CSV — the actual model constraint.
+ # Taking the max rather than the min avoids over-constraining the
+ # condition with the largest legitimate uncertainty.
+ if prior_field is not None:
+ prior_key = _csv_row_name(component, prior_field)
+ prior_updates[prior_key] = float(np.max(clipped_sigma))
+
+ return guess_updates, prior_updates
+
+
def _apply_priors_updates(priors_path, prior_updates):
"""
Overwrite ``parameter == row_name`` rows of the production priors
@@ -691,558 +810,6 @@ def _run_calibration_map(ri,
return svi_state, params, converged
-# ---------------------------------------------------------------------------
-# Diagnostic plots
-# ---------------------------------------------------------------------------
-
-def _make_calibration_plots(gm_cal, params, out_prefix, growth_pred_std=None):
- """
- Generate per-genotype calibration diagnostic plots as PDFs.
-
- For each genotype in the calibration data, writes one PDF containing
- one subplot per (condition_pre, condition_sel, titrant_name, titrant_conc)
- combination. Each subplot shows:
-
- * **Observed** — the experimental ``ln_cfu ± ln_cfu_std`` data points
- plotted at their ``t_sel`` x-coordinates (selection-phase time).
- * **Model prediction** — a smooth 100-point trajectory computed from
- the MAP parameter estimates:
-
- - Pre-selection phase (x from ``-t_pre`` to 0): a straight line from
- ``(−t_pre, ln_cfu0)`` to ``(0, ln_cfu0 + g_pre·t_pre)``, where
- ``ln_cfu0`` comes from the ``ln_cfu0`` deterministic site and the
- pre-selection slope ``g_pre`` is derived from the selection-phase
- intercept.
- - Selection phase (x from 0 to ``max(t_sel)``): a straight line with
- slope ``g_sel`` estimated by fitting the MAP ``growth_pred`` values
- at the observed ``t_sel`` times.
-
- The linear-in-time approximation for the smooth trajectory is exact for
- the ``instant`` growth-transition component and a reasonable visual
- approximation for non-linear variants.
-
- Parameters
- ----------
- gm_cal : ModelOrchestrator
- Calibration ModelOrchestrator (exposes ``growth_tm``, ``data``,
- ``priors``, ``jax_model``).
- params : dict
- MAP parameter dict from the SVI optimiser (keys follow the
- ``{site}_auto_loc`` convention).
- out_prefix : str
- File-name prefix; each genotype's PDF is written to
- ``{out_prefix}_calib_{genotype}.pdf``.
- """
- try:
- import matplotlib
- matplotlib.use("Agg")
- import matplotlib.pyplot as plt
- except ImportError:
- print(
- " warning: matplotlib not available; skipping calibration plots.",
- file=sys.stderr,
- )
- return
-
- from numpyro.infer import Predictive
- from numpyro.infer.autoguide import AutoDelta
- import jax
-
- print("Generating calibration quality plots ...", flush=True)
-
- # --- Run Predictive at the MAP point to recover growth_pred and ln_cfu0 ---
- guide = AutoDelta(gm_cal.jax_model)
- all_indices = jnp.arange(gm_cal.data.num_genotype)
- full_data = gm_cal.get_batch(gm_cal.data, all_indices)
- pred_fn = Predictive(gm_cal.jax_model, guide=guide, params=params, num_samples=1)
- map_samples = pred_fn(
- jax.random.PRNGKey(0), data=full_data, priors=gm_cal.priors
- )
-
- if "growth_pred" not in map_samples:
- print(
- " warning: growth_pred not in model trace; skipping plots.",
- file=sys.stderr,
- )
- return
-
- # Remove leading sample dimension (num_samples=1).
- # growth_pred: (R, T, CP, CS, TN, TC, G) at the original observed timepoints.
- # Kept separately from growth_pred_fine (fine-grid) so the CSV can record
- # predictions that align row-for-row with the observed growth_df entries.
- growth_pred = np.asarray(map_samples["growth_pred"][0])
-
- # ln_cfu0 deterministic: (R, CP, G) — registered by the hierarchical component
- ln_cfu0_map = (
- np.asarray(map_samples["ln_cfu0"][0])
- if "ln_cfu0" in map_samples
- else None
- )
-
- # --- Data tensors from the calibration model ---
- gd = gm_cal.data.growth
- good_mask = np.asarray(gd.good_mask) # (R, T, CP, CS, TN, TC, G)
- t_pre_tensor = np.asarray(gd.t_pre) # (R, T, CP, CS, TN, TC, G)
- t_sel_tensor = np.asarray(gd.t_sel) # (R, T, CP, CS, TN, TC, G)
- ln_cfu_obs = np.asarray(gd.ln_cfu) # (R, T, CP, CS, TN, TC, G)
- ln_cfu_std = np.asarray(gd.ln_cfu_std) # (R, T, CP, CS, TN, TC, G)
-
- n_rep, n_t, n_cp, n_cs, n_tn, n_tc, n_geno = good_mask.shape
-
- # --- Fine-grid Predictive for exact smooth selection-phase trajectories ---
- # Replace the T dimension with T_FINE evenly-spaced points from 0 to the
- # global max t_sel. t_pre is constant over T within each condition cell,
- # so we broadcast from the first observed timepoint. The same pred_fn is
- # reused with the new data argument; JAX will retrace for the new shape.
- T_FINE = 50
- global_max_t_sel = float(np.nanmax(t_sel_tensor[good_mask]))
- t_fine_1d = np.linspace(0.0, global_max_t_sel, T_FINE)
- fine_shape = (n_rep, T_FINE, n_cp, n_cs, n_tn, n_tc, n_geno)
-
- # Helper: broadcast a 7-D array from T=1 (time-constant) to T_FINE.
- # All per-condition tensors are constant along the T axis; slicing any
- # timepoint and broadcasting is therefore exact.
- def _bc_t(arr, *, dtype=None):
- a = np.asarray(arr)
- r = np.broadcast_to(a[:, 0:1, ...], fine_shape).copy()
- return jnp.array(r if dtype is None else r.astype(dtype))
-
- t_sel_fine = np.broadcast_to(
- t_fine_1d[None, :, None, None, None, None, None], fine_shape
- ).copy()
-
- # Active wherever any observed timepoint was valid for that cell.
- has_data_bc = good_mask.any(axis=1, keepdims=True) # (R,1,CP,CS,TN,TC,G)
- good_mask_fine = np.broadcast_to(has_data_bc, fine_shape).copy()
-
- # map_condition_pre / map_condition_sel are also shape (R,T,CP,CS,TN,TC,G);
- # each element is an index into per-condition-rep arrays, constant over T.
- gd_full = full_data.growth
- fine_gd = gd_full.replace(
- num_time=T_FINE,
- t_sel=jnp.array(t_sel_fine),
- t_pre=_bc_t(gd_full.t_pre),
- good_mask=jnp.array(good_mask_fine),
- map_condition_pre=_bc_t(gd_full.map_condition_pre, dtype=int),
- map_condition_sel=_bc_t(gd_full.map_condition_sel, dtype=int),
- ln_cfu=jnp.zeros(fine_shape),
- ln_cfu_std=jnp.ones(fine_shape),
- )
- fine_data = full_data.replace(growth=fine_gd)
- map_samples_fine = pred_fn(
- jax.random.PRNGKey(1), data=fine_data, priors=gm_cal.priors
- )
- # shape: (R, T_FINE, CP, CS, TN, TC, G)
- growth_pred_fine = np.asarray(map_samples_fine["growth_pred"][0])
-
- # --- Dimension labels from the TensorManager ---
- tm = gm_cal.growth_tm
- dn = tm.tensor_dim_names
-
- geno_labels = list(tm.tensor_dim_labels[dn.index("genotype")])
- rep_labels = list(tm.tensor_dim_labels[dn.index("replicate")])
- cp_labels = list(tm.tensor_dim_labels[dn.index("condition_pre")])
- tn_labels = list(tm.tensor_dim_labels[dn.index("titrant_name")])
- tc_labels = list(tm.tensor_dim_labels[dn.index("titrant_conc")])
-
- # Map (cp_idx, cs_idx) → actual condition_sel name. The tensor uses the
- # reduced (integer) condition_sel dimension; the original string name lives
- # in the processed DataFrame alongside the integer index column.
- df = tm.df
- cs_name_map = {}
- for _, row in df.drop_duplicates(
- ["condition_pre_idx", "condition_sel_idx"]
- ).iterrows():
- cs_name_map[
- (int(row["condition_pre_idx"]), int(row["condition_sel_idx"]))
- ] = str(row["condition_sel"])
-
- # Map (genotype, titrant_name, titrant_conc) → mean observed theta.
- # Averages over binding replicates so a single scalar appears in the title.
- theta_map: dict = {}
- for _, row in gm_cal.binding_df.iterrows():
- key = (str(row["genotype"]), str(row["titrant_name"]), float(row["titrant_conc"]))
- theta_map.setdefault(key, []).append(float(row["theta_obs"]))
- theta_map = {k: float(np.nanmean(v)) for k, v in theta_map.items()}
-
- prop_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
-
- for g_i, geno_name in enumerate(geno_labels):
-
- # Collect valid (cp, cs, tn, tc) combinations and the replicates that
- # contribute observations.
- condition_combos: dict = {}
- for r_i in range(n_rep):
- for cp_i in range(n_cp):
- for cs_i in range(n_cs):
- for tn_i in range(n_tn):
- for tc_i in range(n_tc):
- if good_mask[r_i, :, cp_i, cs_i, tn_i, tc_i, g_i].any():
- condition_combos.setdefault(
- (cp_i, cs_i, tn_i, tc_i), []
- ).append(r_i)
-
- if not condition_combos:
- continue
-
- n_combos = len(condition_combos)
- n_cols = min(3, n_combos)
- n_rows = (n_combos + n_cols - 1) // n_cols
-
- fig, axes = plt.subplots(
- n_rows, n_cols,
- figsize=(5 * n_cols, 4 * n_rows),
- squeeze=False,
- sharey=True,
- )
- fig.suptitle(
- f"Calibration fit — genotype: {geno_name}",
- fontsize=13,
- fontweight="bold",
- )
-
- for combo_i, ((cp_i, cs_i, tn_i, tc_i), rep_list) in enumerate(
- condition_combos.items()
- ):
- ax = axes[combo_i // n_cols][combo_i % n_cols]
-
- cp_name = str(cp_labels[cp_i])
- cs_name = cs_name_map.get((cp_i, cs_i), f"sel_{cs_i}")
- tn_name = str(tn_labels[tn_i])
- tc_val = float(tc_labels[tc_i])
-
- theta_obs = theta_map.get((geno_name, tn_name, tc_val))
- theta_str = f", θ = {theta_obs:.3f}" if theta_obs is not None else ""
- ax.set_title(
- f"{cp_name} → {cs_name}\n{tn_name} = {tc_val:.3g}{theta_str}",
- fontsize=9,
- )
- ax.set_xlabel("Time")
- ax.set_ylabel("ln(CFU)")
- ax.axvline(0.0, color="0.6", lw=0.8, ls="--")
-
- # Collect valid per-replicate data for this condition combination.
- rep_data = []
- for r_i in rep_list:
- mask = good_mask[r_i, :, cp_i, cs_i, tn_i, tc_i, g_i]
- valid_t = np.where(mask)[0]
- if len(valid_t) == 0:
- continue
- rep_data.append({
- "r_i": r_i,
- "t_sel": t_sel_tensor[r_i, valid_t, cp_i, cs_i, tn_i, tc_i, g_i],
- "obs": ln_cfu_obs[r_i, valid_t, cp_i, cs_i, tn_i, tc_i, g_i],
- "std": ln_cfu_std[r_i, valid_t, cp_i, cs_i, tn_i, tc_i, g_i],
- "t_pre": float(np.nanmedian(
- t_pre_tensor[r_i, valid_t, cp_i, cs_i, tn_i, tc_i, g_i]
- )),
- })
-
- if not rep_data:
- ax.set_visible(False)
- continue
-
- max_t_sel = float(max(np.nanmax(rd["t_sel"]) for rd in rep_data))
- max_t_pre = float(max(rd["t_pre"] for rd in rep_data))
-
- # Plot each replicate in its own colour: data points + model line.
- for rd in rep_data:
- r_i = rd["r_i"]
- rep_color = prop_colors[r_i % len(prop_colors)]
- rep_label = str(rep_labels[r_i])
-
- # Observed data with error bars
- ax.errorbar(
- rd["t_sel"], rd["obs"], yerr=rd["std"],
- fmt="o", color=rep_color, ms=5, lw=1, capsize=3,
- label=rep_label, zorder=3,
- )
-
- # Exact model prediction for this replicate.
- # growth_pred_fine: (R, T_FINE, CP, CS, TN, TC, G)
- y_fine_r = growth_pred_fine[r_i, :, cp_i, cs_i, tn_i, tc_i, g_i]
- calc_at_0_r = float(y_fine_r[0])
-
- # ln_cfu0 anchor; fall back to calc_at_0 when site is absent.
- if ln_cfu0_map is not None and ln_cfu0_map.ndim == 3:
- ln_cfu0_r = float(ln_cfu0_map[r_i, cp_i, g_i])
- else:
- ln_cfu0_r = calc_at_0_r
-
- # Pre-selection: two-point line (-t_pre, ln_cfu0) → (0, calc_at_0)
- # Selection: exact fine-grid predictions
- t_smooth = np.concatenate([np.array([-rd["t_pre"], 0.0]), t_fine_1d])
- y_smooth = np.concatenate([np.array([ln_cfu0_r, calc_at_0_r]), y_fine_r])
- ax.plot(t_smooth, y_smooth, "-", color=rep_color, lw=1.8, zorder=4)
-
- ax.legend(fontsize=8, loc="best")
- x_pad = (max_t_sel + max_t_pre) * 0.03
- ax.set_xlim(-max_t_pre - x_pad, max_t_sel + x_pad)
-
- # Hide unused subplots
- for extra_i in range(n_combos, n_rows * n_cols):
- axes[extra_i // n_cols][extra_i % n_cols].set_visible(False)
-
- fig.tight_layout(rect=[0, 0, 1, 0.95])
- pdf_path = f"{out_prefix}_calib_{geno_name}.pdf"
- fig.savefig(pdf_path, format="pdf", bbox_inches="tight")
- plt.close(fig)
- print(f" Saved {pdf_path}", flush=True)
-
- # --- Write predictions CSV ---
- # Extract the MAP-predicted ln_cfu at every valid (observed) tensor cell,
- # then attach it as a new column alongside the original growth_df rows.
- #
- # The tensor has 7 dimensions: (R, T, CP, CS, TN, TC, G).
- # tm.df carries _idx columns for 6 of them (all except T); the T dimension
- # is identified by the t_sel float value, which is the same in both the
- # tensor and the DataFrame.
- r_idx, t_idx, cp_idx, cs_idx, tn_idx, tc_idx, g_idx = np.where(good_mask)
-
- idx = (r_idx, t_idx, cp_idx, cs_idx, tn_idx, tc_idx, g_idx)
- pred_lookup_dict = {
- "replicate_idx": r_idx.astype(int),
- "condition_pre_idx": cp_idx.astype(int),
- "condition_sel_idx": cs_idx.astype(int),
- "titrant_name_idx": tn_idx.astype(int),
- "titrant_conc_idx": tc_idx.astype(int),
- "genotype_idx": g_idx.astype(int),
- "t_sel": t_sel_tensor[idx],
- "ln_cfu_pred": growth_pred[idx],
- }
- if growth_pred_std is not None:
- pred_lookup_dict["ln_cfu_pred_std"] = growth_pred_std[idx]
- pred_lookup = pd.DataFrame(pred_lookup_dict)
-
- merge_cols = [
- "replicate_idx", "condition_pre_idx", "condition_sel_idx",
- "titrant_name_idx", "titrant_conc_idx", "genotype_idx", "t_sel",
- ]
- growth_df_out = df.merge(pred_lookup, on=merge_cols, how="left")
- csv_path = f"{out_prefix}_calib_growth_df.csv"
- growth_df_out.to_csv(csv_path, index=False)
- print(f" Saved {csv_path}", flush=True)
-
- n_params = int(sum(np.asarray(v).size
- for k, v in params.items()
- if k.endswith("_auto_loc")))
- n_obs = int(growth_df_out["ln_cfu_pred"].notna().sum())
- _write_calibration_stats(growth_df_out, out_prefix,
- n_params=n_params, n_obs=n_obs)
- _make_correlation_plot(growth_df_out, out_prefix)
-
- print("Calibration plots complete.", flush=True)
-
-
-# ---------------------------------------------------------------------------
-# Calibration stats JSON and correlation plot
-# ---------------------------------------------------------------------------
-
-def _write_calibration_stats(df, out_prefix, n_params=None, n_obs=None):
- """
- Compute goodness-of-fit statistics on the calibration predictions and
- write ``{out_prefix}_calib_stats.json``.
-
- Requires ``ln_cfu``, ``ln_cfu_pred``, and ``ln_cfu_pred_std`` columns in
- ``df``; returns silently if any are absent or if no valid rows remain
- after dropping NaNs.
-
- Parameters
- ----------
- df : pd.DataFrame
- Merged predictions DataFrame.
- out_prefix : str
- File-name prefix for the JSON output.
- n_params : int or None
- Number of free model parameters (from ``_auto_loc`` keys in the MAP
- params dict). Written to JSON as an integer or null.
- n_obs : int or None
- Number of valid observations used for the calibration fit (rows with
- a non-NaN ``ln_cfu_pred``). Written to JSON as an integer or null.
- """
- import json
- from tfscreen.mle import stats_test_suite
-
- needed = ["ln_cfu", "ln_cfu_pred", "ln_cfu_pred_std"]
- if not all(c in df.columns for c in needed):
- return
-
- valid = df.dropna(subset=needed)
- if len(valid) == 0:
- return
-
- stats = stats_test_suite(
- valid["ln_cfu_pred"].to_numpy(dtype=float),
- valid["ln_cfu"].to_numpy(dtype=float),
- valid["ln_cfu_pred_std"].to_numpy(dtype=float),
- )
-
- # json.dump cannot serialise numpy scalars or NaN; normalise to Python
- # float / None so the output is always valid JSON.
- serialisable = {
- k: (None if (v is None or not np.isfinite(float(v))) else float(v))
- for k, v in stats.items()
- }
- serialisable["n_params"] = int(n_params) if n_params is not None else None
- serialisable["n_obs"] = int(n_obs) if n_obs is not None else None
-
- json_path = f"{out_prefix}_calib_stats.json"
- with open(json_path, "w") as fh:
- json.dump(serialisable, fh, indent=2)
- print(f" Saved {json_path}", flush=True)
-
-
-def _make_correlation_plot(df, out_prefix):
- """
- Write ``{out_prefix}_calib_correlation.pdf``: observed vs. predicted
- ln_cfu, coloured by genotype.
-
- Series colours are taken from the active matplotlib prop-cycle, cycling
- for any number of genotypes. Requires ``ln_cfu``, ``ln_cfu_pred``, and
- ``genotype`` columns; returns silently if any are absent.
- """
- try:
- import matplotlib.pyplot as plt
- except ImportError:
- return
-
- from tfscreen.plot.helper import get_ax_limits
-
- needed = ["ln_cfu", "ln_cfu_pred", "genotype"]
- if not all(c in df.columns for c in needed):
- return
-
- valid = df.dropna(subset=["ln_cfu", "ln_cfu_pred"])
- if len(valid) == 0:
- return
-
- genotypes = pd.unique(valid["genotype"])
- prop_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
- colors = {g: prop_colors[i % len(prop_colors)] for i, g in enumerate(genotypes)}
-
- fig, ax = plt.subplots(1, figsize=(6, 6))
- ax.set_aspect("equal")
-
- for g, g_df in valid.groupby("genotype", observed=True):
- ax.scatter(g_df["ln_cfu"], g_df["ln_cfu_pred"],
- s=50, edgecolor=colors[g], facecolor="none", label=g)
-
- lims = get_ax_limits(x_values=valid["ln_cfu"],
- y_values=valid["ln_cfu_pred"],
- pad_by=0.1)
- ax.set_xlim(lims)
- ax.set_ylim(lims)
- ax.plot(lims, lims, "--", color="gray", zorder=-20)
- ax.legend(fontsize=8)
-
- ticks = np.arange(np.ceil(lims[0]), np.floor(lims[1]) + 1, 2)
- ax.set_xticks(ticks)
- ax.set_yticks(ticks)
-
- ax.set_xlabel("ln_cfu obs")
- ax.set_ylabel("ln_cfu pred")
-
- fig.tight_layout()
- pdf_path = f"{out_prefix}_calib_correlation.pdf"
- fig.savefig(pdf_path, format="pdf", bbox_inches="tight")
- plt.close(fig)
- print(f" Saved {pdf_path}", flush=True)
-
-
-# ---------------------------------------------------------------------------
-# Laplace-based prediction uncertainty
-# ---------------------------------------------------------------------------
-
-def _compute_growth_pred_std(ri, params, gm_cal, n_samples=100):
- """
- Estimate per-cell growth_pred standard deviation via Laplace samples.
-
- Recomputes the Hessian of the negative log-joint at the MAP point,
- projects it to the PD cone, draws ``n_samples`` samples from the
- resulting Gaussian, runs the forward model on each, and returns the
- elementwise standard deviation of ``growth_pred`` across samples.
-
- The Hessian computation duplicates work already done in
- ``compute_hessian_sigmas``; this is intentional so that the full
- covariance (not just the diagonal) is available for sampling.
-
- Returns ``None`` if ``growth_pred`` is absent from the model trace or
- if the parameter dict is empty.
- """
- from numpyro.infer.util import potential_energy
- from numpyro.distributions.transforms import biject_to
- from numpyro.infer import Predictive
- import jax.flatten_util
-
- unconstrained = {
- k[: -len("_auto_loc")]: jnp.array(v)
- for k, v in params.items()
- if k.endswith("_auto_loc")
- }
- if not unconstrained:
- return None
-
- data_on_gpu = jax.device_put(gm_cal.data)
- all_indices = jnp.arange(gm_cal.data.num_genotype)
- full_data = gm_cal.get_batch(data_on_gpu, all_indices)
- model_kwargs = {"priors": gm_cal.priors, "data": full_data}
-
- flat_map, unravel = jax.flatten_util.ravel_pytree(unconstrained)
- D = int(flat_map.shape[0])
-
- def pe_fn(flat_p):
- return potential_energy(
- gm_cal.jax_model, [], model_kwargs, unravel(flat_p)
- )
-
- hessian = jax.hessian(pe_fn)(flat_map)
- H_np = np.array(hessian, dtype=np.float64)
- eigenvalues_np, eigenvectors_np = np.linalg.eigh(H_np)
- eigenvalues_pd = np.maximum(eigenvalues_np, 1e-3)
- cov_np = eigenvectors_np @ np.diag(1.0 / eigenvalues_pd) @ eigenvectors_np.T
- L_np = np.linalg.cholesky(cov_np)
- L = jnp.array(L_np, dtype=jnp.float32)
-
- # Per-site unconstrained → constrained bijections.
- model_trace = trace(seed(gm_cal.jax_model, rng_seed=0)).get_trace(**model_kwargs)
- site_transforms = {
- name: biject_to(site["fn"].support)
- for name, site in model_trace.items()
- if site["type"] == "sample" and not site.get("is_observed", False)
- }
-
- # Draw n_samples at once: (n_samples, D).
- sample_key = ri.get_key()
- z = jax.random.normal(sample_key, shape=(n_samples, D))
- flat_samples = flat_map[None] + z @ L.T
-
- unc_samples = jax.vmap(unravel)(flat_samples)
- constrained_samples = {
- k: jax.vmap(site_transforms[k])(v) if k in site_transforms else v
- for k, v in unc_samples.items()
- }
-
- pred_fn = Predictive(gm_cal.jax_model, posterior_samples=constrained_samples)
- pred = pred_fn(ri.get_key(), data=full_data, priors=gm_cal.priors)
-
- if "growth_pred" not in pred:
- return None
-
- # pred["growth_pred"]: (n_samples, R, T, CP, CS, TN, TC, G)
- #
- # Some Laplace samples land far from the MAP (badly-constrained directions
- # have large variance after Hessian eigenvalue clamping). The resulting
- # forward-model predictions can be enormous in float32, and squaring them
- # inside np.std overflows to inf. Two guards:
- # 1. Cast to float64 — extends the safe range from ~3e38 to ~1e308.
- # 2. Clip to ±_GROWTH_PRED_CLIP — any prediction outside this window is
- # physically nonsensical (ln_cfu values live in ~[0, 30]) and would
- # produce a misleadingly large std even in float64.
- _GROWTH_PRED_CLIP = 1e4
- gp = np.asarray(pred["growth_pred"], dtype=np.float64)
- np.clip(gp, -_GROWTH_PRED_CLIP, _GROWTH_PRED_CLIP, out=gp)
- return gp.std(axis=0)
-
# ---------------------------------------------------------------------------
# Public entry point
@@ -1263,7 +830,12 @@ def run_prefit_calibration(config_file,
checkpoint_interval=10,
max_num_epochs=100000,
init_param_jitter=0.0,
- epoch_checkpoint_interval=0):
+ epoch_checkpoint_interval=0,
+ k_scale_floor=_DEFAULT_K_SCALE_FLOOR,
+ m_scale_floor=_DEFAULT_M_SCALE_FLOOR,
+ k_scale_ceiling=_DEFAULT_K_SCALE_CEILING,
+ m_scale_ceiling=_DEFAULT_M_SCALE_CEILING,
+ hessian_chunk_size=64):
"""
Run the calibration pre-fit and update the production priors / guesses
CSVs in place.
@@ -1330,6 +902,30 @@ def run_prefit_calibration(config_file,
Frequency (in epochs) to write numbered epoch checkpoints to a
``checkpoints/`` subdirectory (default 0). Set to 0 or None to
disable.
+ k_scale_floor : float, optional
+ Minimum prior scale applied to the baseline growth-rate (k/b)
+ parameters after Hessian estimation (default
+ ``_DEFAULT_K_SCALE_FLOOR = 0.002``). Represents the irreducible
+ day-to-day variability in growth rates that the Hessian cannot capture.
+ For real data, estimate this by running the prefit on multiple
+ independent calibration replicates and measuring the spread of MAP
+ estimates.
+ m_scale_floor : float, optional
+ Minimum prior scale applied to the theta-slope (m) parameters after
+ Hessian estimation (default ``_DEFAULT_M_SCALE_FLOOR = 0.001``).
+ k_scale_ceiling : float, optional
+ Maximum prior scale for k after Hessian estimation (default
+ ``_DEFAULT_K_SCALE_CEILING = 0.1``, matching the linear component
+ default). Prevents degenerate Hessian directions — which arise from
+ the ``b`` / ``ln_cfu0`` near-degeneracy in the prefit calibration —
+ from producing a scale larger than the original production prior and
+ thereby *loosening* the constraint.
+ m_scale_ceiling : float, optional
+ Maximum prior scale for m_scale_plus after Hessian estimation (default
+ ``_DEFAULT_M_SCALE_CEILING = 0.01``, matching the linear default).
+ hessian_chunk_size : int, optional
+ Number of Hessian rows computed per device batch (default 64).
+ Reduce if device runs out of memory during Hessian computation.
Returns
-------
@@ -1345,11 +941,11 @@ def run_prefit_calibration(config_file,
# 1. Resolve production config and CSV targets.
priors_path, guesses_path = _resolve_csv_paths(config_file)
- gm_prod, _ = read_configuration(config_file)
+ orchestrator_prod, _ = read_configuration(config_file)
# 2. Filter to the calibration (genotype, titrant_name, titrant_conc)
# intersection. read_configuration already loaded the production
- # data paths into gm_prod; pull them straight off the config so we
+ # data paths into orchestrator_prod; pull them straight off the config so we
# don't depend on an internal attribute.
with open(config_file, "r") as fh:
cfg = yaml.safe_load(fh)
@@ -1358,17 +954,17 @@ def run_prefit_calibration(config_file,
growth_df_cal, binding_df_cal = _intersect_data(growth_path, binding_path)
# 3. Build the calibration model with overrides applied.
- gm_cal = _build_calibration_model(gm_prod, growth_df_cal, binding_df_cal)
- theta_values = _compute_theta_values(gm_cal, binding_df_cal)
- _inject_calibration_priors(gm_cal, gm_prod, theta_values)
+ orchestrator_cal = _build_calibration_model(orchestrator_prod, growth_df_cal, binding_df_cal)
+ theta_values = _compute_theta_values(orchestrator_cal, binding_df_cal)
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values)
# 4. MAP fit + Hessian sigmas.
effective_seed = seed if seed is not None else 0
- ri = RunInference(gm_cal, effective_seed)
+ ri = RunInference(orchestrator_cal, effective_seed)
svi_state, params, converged = _run_calibration_map(
ri,
- init_params=gm_cal.init_params,
+ init_params=orchestrator_cal.init_params,
out_prefix=out_prefix,
checkpoint_file=checkpoint_file,
adam_step_size=adam_step_size,
@@ -1385,21 +981,30 @@ def run_prefit_calibration(config_file,
epoch_checkpoint_interval=epoch_checkpoint_interval
)
- print("Computing Hessian-based per-site uncertainties ...", flush=True)
- hessian_results = ri.compute_hessian_sigmas(params)
+ # 5. Hessian-derived scale estimation.
+ # Compute the diagonal of the inverse Hessian at the MAP to get per-
+ # parameter uncertainty estimates. These are floored and written as
+ # per-condition scales (guesses) and scalar prior scales (priors) so the
+ # production SVI cannot drift far from the calibrated growth parameters.
+ hessian_results = ri.compute_hessian_sigmas(
+ params, hessian_chunk_size=hessian_chunk_size
+ )
+
+ # 6. Map sample sites → CSV fields and apply in-place updates.
+ field_mapping = _identify_field_mapping(orchestrator_cal)
+ prior_updates, guess_updates = _build_csv_updates(field_mapping, params)
+
+ hessian_guess_updates, hessian_prior_updates = _build_hessian_scale_updates(
+ field_mapping, hessian_results,
+ k_scale_floor, m_scale_floor,
+ k_scale_ceiling, m_scale_ceiling,
+ )
+ guess_updates.update(hessian_guess_updates)
+ prior_updates.update(hessian_prior_updates)
- # 5. Map sample sites → CSV fields and apply in-place updates.
- field_mapping = _identify_field_mapping(gm_cal)
- prior_updates, guess_updates = _build_csv_updates(field_mapping,
- hessian_results)
_apply_priors_updates(priors_path, prior_updates)
_apply_guesses_updates(guesses_path, guess_updates)
- # 6. Write per-genotype diagnostic plots with Laplace-based prediction uncertainty.
- print("Computing growth_pred uncertainty via Laplace samples ...", flush=True)
- growth_pred_std = _compute_growth_pred_std(ri, params, gm_cal)
- _make_calibration_plots(gm_cal, params, out_prefix, growth_pred_std=growth_pred_std)
-
return svi_state, params, converged
@@ -1409,7 +1014,12 @@ def main():
manual_arg_types={"config_file": str,
"seed": int,
"checkpoint_file": str,
- "init_param_jitter": float},
+ "init_param_jitter": float,
+ "k_scale_floor": float,
+ "m_scale_floor": float,
+ "k_scale_ceiling": float,
+ "m_scale_ceiling": float,
+ "hessian_chunk_size": int},
)
diff --git a/src/tfscreen/tfmodel/scripts/sample_posterior_cli.py b/src/tfscreen/tfmodel/scripts/sample_posterior_cli.py
index 1c9e7555..11ef71bf 100644
--- a/src/tfscreen/tfmodel/scripts/sample_posterior_cli.py
+++ b/src/tfscreen/tfmodel/scripts/sample_posterior_cli.py
@@ -2,7 +2,6 @@
import dill
from tfscreen.tfmodel.configuration_io import read_configuration
from tfscreen.tfmodel.inference.run_inference import RunInference
-from tfscreen.tfmodel.scripts.fit_model_cli import _run_svi
from tfscreen.util.cli.generalized_main import generalized_main
@@ -23,8 +22,8 @@ def sample_posterior(config_file,
writes posterior predictives.
2. MAP checkpoint (AutoDelta guide): forms a Laplace (Gaussian) approximation
via the Hessian of the log-joint at the MAP point, then draws samples.
- 3. SVI checkpoint (component guide): resumes the fitted guide with 0 additional
- epochs and draws posterior samples directly.
+ 3. SVI checkpoint (component guide): restores the fitted variational state
+ and draws posterior samples directly from the guide.
The .h5 file written by this command can be passed directly to
tfs-predict-growth, tfs-predict-theta, and tfs-extract-params to obtain
@@ -63,8 +62,8 @@ def sample_posterior(config_file,
"Run tfs-fit-model first to produce a checkpoint."
)
- gm, init_params = read_configuration(config_file)
- ri = RunInference(gm, seed)
+ orchestrator, init_params = read_configuration(config_file)
+ ri = RunInference(orchestrator, seed)
with open(checkpoint_file, "rb") as f:
chk_data = dill.load(f)
@@ -95,30 +94,18 @@ def sample_posterior(config_file,
hessian_chunk_size=hessian_chunk_size,
)
else:
- # SVI checkpoint: resume with 0 epochs, draw samples directly.
+ # SVI checkpoint: rebuild the guide object then restore the saved
+ # variational state directly — no optimization loop needed.
print("Detected SVI checkpoint. Drawing variational posterior samples...", flush=True)
- _run_svi(ri,
- init_params=init_params,
- checkpoint_file=checkpoint_file,
- out_prefix=ri_prefix,
- max_num_epochs=0,
- num_posterior_samples=num_posterior_samples,
- sampling_batch_size=sampling_batch_size,
- forward_batch_size=forward_batch_size,
- always_get_posterior=True,
- # Convergence / optimizer kwargs are unused at 0 epochs but
- # required by _run_svi's signature; use neutral defaults.
- adam_step_size=1e-3,
- adam_final_step_size=1e-6,
- adam_clip_norm=1.0,
- elbo_num_particles=2,
- convergence_tolerance=1e-5,
- convergence_window=10,
- patience=10,
- convergence_check_interval=2,
- checkpoint_interval=10,
- init_param_jitter=0.0,
- epoch_checkpoint_interval=None)
+ svi_obj, svi_state = ri.restore_svi_from_checkpoint(
+ checkpoint_file, init_params=init_params
+ )
+ ri.get_posteriors(svi=svi_obj,
+ svi_state=svi_state,
+ out_prefix=ri_prefix,
+ num_posterior_samples=num_posterior_samples,
+ sampling_batch_size=sampling_batch_size,
+ forward_batch_size=forward_batch_size)
src = f"{ri_prefix}_posterior.h5"
dst = f"{out_prefix}.h5"
diff --git a/src/tfscreen/tfmodel/scripts/prior_predictive_cli.py b/src/tfscreen/tfmodel/scripts/sample_prior_cli.py
similarity index 93%
rename from src/tfscreen/tfmodel/scripts/prior_predictive_cli.py
rename to src/tfscreen/tfmodel/scripts/sample_prior_cli.py
index 94723de7..3572ee4d 100644
--- a/src/tfscreen/tfmodel/scripts/prior_predictive_cli.py
+++ b/src/tfscreen/tfmodel/scripts/sample_prior_cli.py
@@ -67,7 +67,7 @@ def sample_prior(config_file,
independence. Default 0.
"""
print(f"Loading configuration from {config_file}...", flush=True)
- gm, _ = read_configuration(config_file)
+ orchestrator, _ = read_configuration(config_file)
width = len(str(num_datasets))
@@ -75,10 +75,10 @@ def sample_prior(config_file,
tag = str(i).zfill(max(width, 3))
print(f"Drawing prior sample {i + 1}/{num_datasets}...", flush=True)
- predictions, latent_params = draw_prior(gm, rng_key=seed + i, num_draws=1)
+ predictions, latent_params = draw_prior(orchestrator, rng_key=seed + i, num_draws=1)
rng = np.random.default_rng(seed + i) if noise else None
- growth_df = growth_df_from_prior(gm, latent_params, draw_idx=0, noise_rng=rng)
+ growth_df = growth_df_from_prior(orchestrator, latent_params, draw_idx=0, noise_rng=rng)
csv_path = f"{out_prefix}_{tag}_growth.csv"
growth_df.to_csv(csv_path, index=False)
diff --git a/src/tfscreen/tfmodel/scripts/subset_genotypes_cli.py b/src/tfscreen/tfmodel/scripts/subset_genotypes_cli.py
index 23f8e74a..2da39f26 100644
--- a/src/tfscreen/tfmodel/scripts/subset_genotypes_cli.py
+++ b/src/tfscreen/tfmodel/scripts/subset_genotypes_cli.py
@@ -41,7 +41,7 @@ def subset_genotypes(
data_file,
n_singles,
n_steps,
- out_prefix="genotype_subset",
+ out_prefix="tfs_subset",
whitelist_file=None,
blacklist_file=None,
random_seed=42,
@@ -70,7 +70,7 @@ def subset_genotypes(
included in training are spaced linearly from 0 to all possible,
with duplicate counts collapsed.
out_prefix : str, optional
- Prefix for all output files. Default 'genotype_subset'.
+ Prefix for all output files. Default 'tfs_subset'.
whitelist_file : str, optional
Path to a plain-text file (one genotype per line, # comments allowed)
of genotypes that must be included in the selected singles.
diff --git a/src/tfscreen/tfmodel/scripts/summarize_fit_cli.py b/src/tfscreen/tfmodel/scripts/summarize_fit_cli.py
index f210e942..c7e4d88a 100644
--- a/src/tfscreen/tfmodel/scripts/summarize_fit_cli.py
+++ b/src/tfscreen/tfmodel/scripts/summarize_fit_cli.py
@@ -2,6 +2,7 @@
import glob
import json
import os
+import re
import warnings
import numpy as np
@@ -9,11 +10,121 @@
import yaml
from matplotlib import pyplot as plt
-from tfscreen.mle.stats_test_suite import stats_test_suite
+from tfscreen.analysis.stats_test_suite import stats_test_suite
+from tfscreen.plot import default_styles # noqa: F401 — applies global rcParams on import
+from tfscreen.plot.default_styles import DEFAULT_COLORS
from tfscreen.plot.xy_corr import xy_corr
+from tfscreen.tfmodel.analysis.error_calibration import calibration_summary
from tfscreen.util.cli.generalized_main import generalized_main
+def _find_params_or_posterior(run_dir):
+ """Return (kind, path) for the best available prediction source in run_dir.
+
+ Preference order: ``*_posterior.h5`` first (richer uncertainty), then
+ ``*_params.npz`` (MAP fallback). Returns ``(None, None)`` when neither
+ is found. Warns when multiple files of the same type are present and
+ uses the first alphabetically.
+ """
+ for suffix, kind in (("_posterior.h5", "posterior"), ("_params.npz", "params")):
+ matches = sorted(glob.glob(os.path.join(run_dir, f"*{suffix}")))
+ if not matches:
+ continue
+ if len(matches) > 1:
+ warnings.warn(
+ f"Multiple {kind} files found in {run_dir}; "
+ f"using {os.path.basename(matches[0])}"
+ )
+ return kind, matches[0]
+ return None, None
+
+
+def _try_plot_theta_fits(binding_df, pred_df, out_prefix):
+ """Generate per-genotype theta fit plots merging binding observations with predictions.
+
+ Silently skips when binding_df or pred_df is None. Warns and returns on
+ any failure so that the rest of summarize_fit output is unaffected.
+ """
+ if binding_df is None or pred_df is None:
+ return
+
+ if "theta_std" not in binding_df.columns:
+ warnings.warn(
+ "binding_df has no 'theta_std' column; skipping theta fit plots."
+ )
+ return
+
+ try:
+ from tfscreen.plot.plot_theta_fits import plot_theta_fits
+
+ join_cols = ["genotype", "titrant_name", "titrant_conc"]
+ binding_sel = [c for c in join_cols + ["theta_obs", "theta_std"]
+ if c in binding_df.columns]
+ merged = pred_df.merge(binding_df[binding_sel], on=join_cols, how="inner")
+ if len(merged) == 0:
+ return
+
+ for geno in sorted(merged["genotype"].unique().tolist(), key=str):
+ geno_df = merged[merged["genotype"] == geno]
+ safe_name = str(geno).replace("/", "_").replace(" ", "_")
+ csv_path = f"{out_prefix}_{safe_name}_theta_fits.csv"
+ geno_df.to_csv(csv_path, index=False)
+ print(f"Wrote theta fit data to {csv_path}")
+ ax = plot_theta_fits(geno_df, title=str(geno))
+ fig = ax.get_figure()
+ pdf_path = f"{out_prefix}_{safe_name}_theta_fits.pdf"
+ fig.savefig(pdf_path, format="pdf", bbox_inches="tight")
+ plt.close(fig)
+ print(f"Wrote theta fit plot to {pdf_path}")
+ except Exception as exc:
+ warnings.warn(f"Could not generate theta fit plots: {exc}")
+
+
+def _try_plot_trajectories(config_file, config_yaml, run_dir, out_prefix, binding_df):
+ """Attempt to generate per-genotype growth trajectory plots.
+
+ Silently skips when the run has no growth data. Warns and returns on any
+ other failure so that the rest of summarize_fit output is unaffected.
+ """
+ if config_yaml is None or not config_yaml.get("data", {}).get("growth"):
+ return
+
+ kind, pred_path = _find_params_or_posterior(run_dir)
+ if pred_path is None:
+ warnings.warn(
+ "No *_posterior.h5 or *_params.npz found in "
+ f"{run_dir}; skipping trajectory plots."
+ )
+ return
+
+ try:
+ from tfscreen.tfmodel.configuration_io import read_configuration
+ from tfscreen.plot.geno_trajectory import predict_geno_trajectory_df, plot_geno_trajectory
+
+ orchestrator, _ = read_configuration(config_file)
+
+ genotypes = list(binding_df["genotype"].unique()) if binding_df is not None else None
+
+ pred_df = predict_geno_trajectory_df(
+ orchestrator,
+ pred_path,
+ genotypes=genotypes,
+ )
+ for geno in sorted(pred_df["genotype"].unique().tolist(), key=str):
+ geno_df = pred_df[pred_df["genotype"] == geno]
+ safe_name = str(geno).replace("/", "_").replace(" ", "_")
+ csv_path = f"{out_prefix}_{safe_name}_trajectory.csv"
+ geno_df.to_csv(csv_path, index=False)
+ print(f"Wrote trajectory data to {csv_path}")
+ fig = plot_geno_trajectory(geno_df)
+ pdf_path = f"{out_prefix}_{safe_name}_trajectory.pdf"
+ fig.savefig(pdf_path, format="pdf", bbox_inches="tight")
+ plt.close(fig)
+ print(f"Wrote trajectory plot to {pdf_path}")
+ except Exception as exc:
+ warnings.warn(f"Could not generate trajectory plots: {exc}")
+
+
def _find_unique(run_dir, suffix, label, warn_missing=True):
"""Return single file in run_dir matching *suffix, or None."""
matches = sorted(glob.glob(os.path.join(run_dir, f"*{suffix}")))
@@ -122,23 +233,328 @@ def _blank_panel(ax, message):
ha="center", va="center", fontsize=12, color="gray")
+def _extract_quantile_data(df):
+ """Return (quantile_matrix, levels) from a DataFrame with q{level} columns.
+
+ Columns must match ``q{float}`` exactly (e.g. ``q0.025``, ``q0.5``).
+ Returns ``(None, None)`` when fewer than two such columns exist.
+ """
+ q_cols = sorted(
+ [c for c in df.columns if re.match(r"^q\d*\.?\d+$", c)],
+ key=lambda c: float(c[1:]),
+ )
+ if len(q_cols) < 2:
+ return None, None
+ levels = np.array([float(c[1:]) for c in q_cols])
+ return df[q_cols].values, levels
+
+
+def _try_calibration(df, ref_col, out_prefix, label):
+ """Run calibration_summary on df[ref_col] vs quantile columns.
+
+ Silently skips when df is None, ref_col is absent, or fewer than two
+ quantile columns exist. Warns and returns on any other failure so the
+ rest of the output is unaffected.
+ """
+ if df is None:
+ return
+ if ref_col not in df.columns:
+ return
+ quantile_matrix, levels = _extract_quantile_data(df)
+ if quantile_matrix is None:
+ return
+ try:
+ calibration_summary(
+ true_vals=df[ref_col].values,
+ quantile_matrix=quantile_matrix,
+ quantile_levels=levels,
+ out_prefix=out_prefix,
+ label=label,
+ )
+ except Exception as exc:
+ warnings.warn(f"Could not run calibration summary for {label}: {exc}")
+
+
+def _build_transform_registry(sim_df):
+ """Return a dict mapping parameter name → ref Series indexed by genotype.
+
+ Supports four resolution strategies (tried in this order at registration):
+
+ 1. **Compound**: ``logit_delta`` = logit(theta_high) – logit(theta_low),
+ registered before the column loop so it is never overwritten.
+ 2. **Direct**: column name is the parameter name.
+ 3. **log_ prefix**: ``log_{col}`` → ``np.log(sim[col])``.
+ 4. **logit_ prefix**: ``logit_{col}`` → logit(sim[col]). For columns
+ that start with ``theta_`` the short form is also registered, e.g.
+ ``theta_low`` registers both ``logit_theta_low`` and ``logit_low``.
+
+ ``setdefault`` is used for entries 2–4 so compound entries win.
+ """
+ sim_lookup = sim_df.set_index("genotype")
+ sim_cols = [c for c in sim_df.columns if c != "genotype"]
+ registry = {}
+
+ # Compound: logit_delta = logit(theta_high) - logit(theta_low)
+ if "theta_high" in sim_cols and "theta_low" in sim_cols:
+ high = sim_lookup["theta_high"].clip(1e-9, 1 - 1e-9)
+ low = sim_lookup["theta_low"].clip(1e-9, 1 - 1e-9)
+ registry["logit_delta"] = np.log(high / (1 - high)) - np.log(low / (1 - low))
+
+ for col in sim_cols:
+ col_s = sim_lookup[col]
+ clipped = col_s.clip(1e-9, 1 - 1e-9)
+ logit_s = np.log(clipped / (1 - clipped))
+
+ registry.setdefault(col, col_s)
+ registry.setdefault(f"log_{col}", np.log(col_s))
+ registry.setdefault(f"logit_{col}", logit_s)
+
+ # For theta_* columns also register the short logit/log/direct forms,
+ # e.g. theta_low → logit_low, log_low, low.
+ if col.startswith("theta_"):
+ base = col[6:]
+ registry.setdefault(f"logit_{base}", logit_s)
+ registry.setdefault(f"log_{base}", np.log(col_s))
+ registry.setdefault(base, col_s)
+
+ return registry
+
+
+def _compute_direct_obs(params_df, ref_series):
+ """Return ref array for a direct params file keyed on genotype.
+
+ Returns ``None`` when params_df has no ``genotype`` column (e.g. files
+ keyed on condition/replicate that cannot be matched to sim_parameters).
+ """
+ if "genotype" not in params_df.columns:
+ return None
+ return ref_series.reindex(params_df["genotype"]).values
+
+
+def _compute_diff_obs(params_df, ref_series):
+ """Return ref array for a diff params file keyed on mutation.
+
+ ref = sim[mutation] − sim[wt]. Returns ``None`` when params_df has no
+ ``mutation`` column. Sets ref to NaN when ``wt`` is absent from
+ ref_series.
+ """
+ if "mutation" not in params_df.columns:
+ return None
+ if "wt" not in ref_series.index:
+ warnings.warn("'wt' not found in sim_parameters; setting diff ref to NaN")
+ return np.full(len(params_df), np.nan)
+ wt_val = ref_series["wt"]
+ return ref_series.reindex(params_df["mutation"]).values - wt_val
+
+
+def _compute_epi_obs(params_df, ref_series, param_name):
+ """Return ref array for an epi params file keyed on pair.
+
+ Uses ``extract_epistasis`` (additive scale) to compute per-double-mutant
+ epistasis from the ground-truth ref_series. Genotypes not present in the
+ result (triple+, unrecognised pairs) receive NaN. Returns ``None`` when
+ params_df has no ``pair`` column.
+ """
+ from tfscreen.analysis.extract_epistasis import extract_epistasis
+
+ if "pair" not in params_df.columns:
+ return None
+
+ work_df = pd.DataFrame({
+ "genotype": ref_series.index,
+ "_obs": ref_series.values,
+ })
+ try:
+ ep_result = extract_epistasis(work_df, y_obs="_obs", scale="add")
+ except Exception as exc:
+ warnings.warn(f"extract_epistasis failed for {param_name}: {exc}")
+ return np.full(len(params_df), np.nan)
+
+ if len(ep_result) == 0:
+ return np.full(len(params_df), np.nan)
+
+ from tfscreen.genetics import standardize_genotypes
+
+ ep_map = ep_result.set_index("genotype")["ep_obs"]
+ std_pairs = pd.Series(standardize_genotypes(params_df["pair"]),
+ index=params_df.index)
+ return ep_map.reindex(std_pairs).values
+
+
+def _try_plot_params_corr(out_df, label, pdf_path):
+ """Save a scatter correlation plot of ref vs median for a params CSV.
+
+ Silently skips when fewer than two finite (ref, median) pairs exist.
+ Warns and returns on any other failure so the rest of the output is
+ unaffected.
+ """
+ try:
+ valid = out_df[["ref", "q0.5"]].dropna()
+ if len(valid) < 2:
+ return
+ fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+ xy_corr(
+ x_values=valid["ref"].values,
+ y_values=valid["q0.5"].values,
+ as_hexbin=False,
+ ax=ax,
+ )
+ ax.set_xlabel(f"Simulated {label}")
+ ax.set_ylabel(f"Predicted {label}")
+ ax.set_title(label)
+ fig.tight_layout()
+ fig.savefig(pdf_path, format="pdf", bbox_inches="tight")
+ plt.close(fig)
+ print(f"Wrote parameter correlation plot to {pdf_path}")
+ except Exception as exc:
+ warnings.warn(f"Could not generate parameter correlation plot {pdf_path}: {exc}")
+
+
+def _summarize_params(run_dir, out_prefix):
+ """Annotate *_params_*.csv files with ground-truth ref from tfs_sim_parameters.csv.
+
+ Scans run_dir for a ``*_sim_parameters.csv`` file. If absent the function
+ returns silently (real-data runs have no ground truth).
+
+ For each non-genotype column in sim_parameters the function checks that a
+ matching ``*_params_{col}.csv`` exists and warns if it does not.
+
+ All ``*_params_*.csv`` files in run_dir are then classified by their
+ suffix as *direct* (genotype key), *diff* (``_d_`` infix, mutation key),
+ or *epi* (``_epi_`` infix, pair key). A ``ref`` column is appended to
+ each file whose parameter name resolves via the transform registry and
+ written to the summary directory under ``{out_prefix}_params_*.csv``.
+ """
+ sim_path = _find_unique(run_dir, "_sim_parameters.csv", "sim_parameters",
+ warn_missing=False)
+ if sim_path is None:
+ return
+
+ try:
+ sim_df = pd.read_csv(sim_path)
+ except Exception as exc:
+ warnings.warn(f"Could not load sim parameters from {sim_path}: {exc}")
+ return
+
+ if "genotype" not in sim_df.columns:
+ warnings.warn(
+ f"sim_parameters {sim_path} has no 'genotype' column; "
+ "skipping param summaries"
+ )
+ return
+
+ sim_cols = [c for c in sim_df.columns if c != "genotype"]
+
+ for col in sim_cols:
+ if not glob.glob(os.path.join(run_dir, f"*_params_{col}.csv")):
+ warnings.warn(f"Expected *_params_{col}.csv in {run_dir} but none found")
+
+ transform_registry = _build_transform_registry(sim_df)
+
+ for param_file in sorted(glob.glob(os.path.join(run_dir, "*_params_*.csv"))):
+ basename = os.path.basename(param_file)
+
+ m_epi = re.search(r"_params_epi_(.+)\.csv$", basename)
+ m_diff = re.search(r"_params_d_(.+)\.csv$", basename)
+ m_direct = re.search(r"_params_(.+)\.csv$", basename)
+
+ if m_epi:
+ kind, param_name = "epi", m_epi.group(1)
+ elif m_diff:
+ kind, param_name = "diff", m_diff.group(1)
+ elif m_direct:
+ kind, param_name = "direct", m_direct.group(1)
+ else:
+ continue
+
+ if param_name not in transform_registry:
+ continue
+
+ ref_series = transform_registry[param_name]
+
+ try:
+ params_df = pd.read_csv(param_file)
+ except Exception as exc:
+ warnings.warn(f"Could not read {param_file}: {exc}")
+ continue
+
+ ref = None
+ try:
+ if kind == "direct":
+ ref = _compute_direct_obs(params_df, ref_series)
+ elif kind == "diff":
+ ref = _compute_diff_obs(params_df, ref_series)
+ else:
+ ref = _compute_epi_obs(params_df, ref_series, param_name)
+ except Exception as exc:
+ warnings.warn(f"Could not compute ref for {basename}: {exc}")
+ continue
+
+ if ref is None:
+ continue
+
+ out_df = params_df.copy()
+ out_df["ref"] = ref
+
+ if kind == "epi":
+ out_name = f"{out_prefix}_params_epi_{param_name}.csv"
+ plot_label = f"epi_{param_name}"
+ elif kind == "diff":
+ out_name = f"{out_prefix}_params_d_{param_name}.csv"
+ plot_label = f"d_{param_name}"
+ else:
+ out_name = f"{out_prefix}_params_{param_name}.csv"
+ plot_label = param_name
+
+ try:
+ out_df.to_csv(out_name, index=False)
+ print(f"Wrote parameter summary to {out_name}")
+ except Exception as exc:
+ warnings.warn(f"Could not write {out_name}: {exc}")
+ continue
+
+ pdf_name = out_name.replace(".csv", ".pdf")
+ _try_plot_params_corr(out_df, plot_label, pdf_name)
+ _try_calibration(out_df, "ref", out_name[:-4], plot_label)
+
+
def summarize_fit(run_dir,
- ground_truth_file=None,
+ ref_theta_file=None,
out_prefix=None):
"""
Evaluate theta prediction quality for a tfs-fit-model run directory.
- Scans run_dir for a *_config.yaml, *_theta_pred.csv, *_growth_pred.csv
+ Scans run_dir for a *_config.yaml, *_pred_theta.csv, *_pred_growth.csv
(optional), and *_losses.txt (optional). Computes prediction statistics
- and writes three output files:
+ and writes output files to the summary directory.
+
+ **CSV files are the authoritative outputs.** Every plot written by this
+ function has a matched CSV containing the exact data used to generate it.
+ The PDFs are human-readable summaries; the CSVs are the record of truth
+ for downstream quantitative analysis.
+
+ Output files written:
- ``{out_prefix}_fit_summary.json`` — nested statistics with top-level
keys ``metadata``, ``theta`` (training / test sub-keys), and ``growth``
(training sub-key).
- - ``{out_prefix}_theta_corr.pdf`` — two-panel hexbin correlation plot for
- theta (training left, test right).
- - ``{out_prefix}_growth_corr.pdf`` — hexbin correlation plot for ln_cfu
- (only written when *_growth_pred.csv is present).
+ - ``{out_prefix}_theta_corr_training.csv`` — joined (ref, predicted)
+ theta pairs for training genotypes; ``ref`` column holds the training
+ observation used as ground truth.
+ - ``{out_prefix}_theta_corr_test.csv`` — same for test genotypes (only
+ written when a ref theta file is resolved and has matching rows).
+ - ``{out_prefix}_theta_corr.pdf`` — two-panel correlation plot for theta
+ (training left, test right).
+ - ``{out_prefix}_growth_corr.csv`` — relative symlink to the
+ *_pred_growth.csv in run_dir (only created when that file is present).
+ - ``{out_prefix}_growth_corr.pdf`` — correlation plot for ln_cfu (only
+ written when *_pred_growth.csv is present).
+ - ``{out_prefix}_{genotype}_theta_fits.csv`` / ``.pdf`` — per-genotype
+ joined observation + prediction data and fit plot (one pair per
+ genotype present in both binding and prediction CSVs).
+ - ``{out_prefix}_{genotype}_trajectory.csv`` / ``.pdf`` — per-genotype
+ predicted growth trajectory data and plot (only when growth data is
+ configured and a posterior/params file is found).
- ``{out_prefix}_losses.pdf`` — training loss curve (only written when
*_losses.txt is present).
@@ -146,20 +562,28 @@ def summarize_fit(run_dir,
----------
run_dir : str
Directory containing model fit outputs.
- ground_truth_file : str, optional
- CSV with columns genotype, titrant_name, titrant_conc, theta_obs,
- theta_std containing known theta values for evaluating out-of-sample
- prediction. When omitted, only training statistics are computed.
+ ref_theta_file : str, optional
+ CSV with columns genotype, titrant_name, titrant_conc, and a theta
+ column containing known theta values for evaluating out-of-sample
+ prediction. When omitted the function looks for a
+ ``*_sim_genotype_theta.csv`` file in *run_dir* automatically (the
+ file produced by ``tfs-simulate``). The theta column is resolved
+ in order: ``theta_obs`` (preferred) → ``theta`` (fallback); it is
+ renamed to ``ref`` in the output CSV. If neither column is present
+ a warning is issued and test statistics are skipped. Pass this
+ argument explicitly to use a theta reference from outside *run_dir*
+ or with a non-standard name.
out_prefix : str, optional
Prefix for output files. Defaults to {run_dir}/tfs_summarize.
"""
run_dir = os.path.abspath(run_dir)
if out_prefix is None:
- out_prefix = os.path.join(run_dir, "tfs_summarize")
+ out_prefix = os.path.join(run_dir, "summary", "tfs_summarize")
+ os.makedirs(os.path.dirname(os.path.abspath(out_prefix)), exist_ok=True)
metadata = {
"run_dir": run_dir,
- "ground_truth_file": ground_truth_file,
+ "ref_theta_file": None,
"timestamp": datetime.datetime.now().isoformat(),
"n_parameters": None,
"n_theta_training_points": None,
@@ -176,9 +600,9 @@ def summarize_fit(run_dir,
# --- Locate files in run_dir ---
config_file = _find_unique(run_dir, "_config.yaml", "config")
- theta_pred_file = _find_unique(run_dir, "_theta_pred.csv", "theta predictions")
+ theta_pred_file = _find_unique(run_dir, "_pred_theta.csv", "theta predictions")
losses_file = _find_unique(run_dir, "_losses.txt", "losses", warn_missing=False)
- growth_pred_file = _find_unique(run_dir, "_growth_pred.csv", "growth predictions",
+ growth_pred_file = _find_unique(run_dir, "_pred_growth.csv", "growth predictions",
warn_missing=False)
# --- Parse YAML config once ---
@@ -246,50 +670,83 @@ def summarize_fit(run_dir,
binding_df[join_cols + ["theta_obs"]],
on=join_cols,
how="inner",
- )
+ ).rename(columns={"theta_obs": "ref"})
metadata["n_theta_training_points"] = len(train_merged)
if len(train_merged) > 0:
theta_training_stats = _run_stats(
- train_merged["median"].values,
- train_merged["theta_obs"].values,
+ train_merged["q0.5"].values,
+ train_merged["ref"].values,
+ )
+ _try_calibration(
+ train_merged, "ref",
+ f"{out_prefix}_theta_training", "theta training",
)
except Exception as exc:
warnings.warn(f"Could not compute theta training statistics: {exc}")
+ # --- Resolve ref_theta_file: explicit arg wins, else auto-discover ---
+ if ref_theta_file is None:
+ ref_theta_file = _find_unique(run_dir, "_sim_genotype_theta.csv",
+ "ref theta", warn_missing=False)
+ metadata["ref_theta_file"] = ref_theta_file
+
# --- Theta test statistics ---
- if ground_truth_file is not None and pred_df is not None:
+ if ref_theta_file is not None and pred_df is not None:
try:
- gt_df = pd.read_csv(ground_truth_file)
+ gt_df = pd.read_csv(ref_theta_file)
join_cols = ["genotype", "titrant_name", "titrant_conc"]
+ if "theta_obs" in gt_df.columns:
+ theta_col = "theta_obs"
+ elif "theta" in gt_df.columns:
+ theta_col = "theta"
+ else:
+ warnings.warn(
+ f"Ref theta file {ref_theta_file} has neither "
+ f"'theta_obs' nor 'theta' column; skipping test statistics"
+ )
+ theta_col = None
+ if theta_col is not None:
+ gt_df = gt_df.rename(columns={theta_col: "ref"})
test_merged = pred_df.merge(
- gt_df[join_cols + ["theta_obs"]],
+ gt_df[join_cols + ["ref"]],
on=join_cols,
how="inner",
- )
- metadata["n_theta_test_points"] = len(test_merged)
- if len(test_merged) > 0:
- theta_test_stats = _run_stats(
- test_merged["median"].values,
- test_merged["theta_obs"].values,
- )
+ ) if theta_col is not None else None
+ if test_merged is not None:
+ metadata["n_theta_test_points"] = len(test_merged)
+ if len(test_merged) > 0:
+ theta_test_stats = _run_stats(
+ test_merged["q0.5"].values,
+ test_merged["ref"].values,
+ )
+ _try_calibration(
+ test_merged, "ref",
+ f"{out_prefix}_theta_test", "theta test",
+ )
except Exception as exc:
- warnings.warn(f"Could not compute theta test statistics: {exc}")
+ warnings.warn(f"Could not compute theta test statistics from {ref_theta_file}: {exc}")
# --- Growth training statistics ---
if growth_pred_file is not None:
try:
growth_pred_df = pd.read_csv(growth_pred_file)
# Drop rows where observed ln_cfu is missing (no observation at that timepoint)
- growth_valid = growth_pred_df.dropna(subset=["ln_cfu", "median"])
+ growth_valid = growth_pred_df.dropna(subset=["ln_cfu", "q0.5"])
metadata["n_growth_training_points"] = len(growth_valid)
if len(growth_valid) > 0:
growth_training_stats = _run_stats(
- growth_valid["median"].values,
+ growth_valid["q0.5"].values,
growth_valid["ln_cfu"].values,
)
except Exception as exc:
warnings.warn(f"Could not compute growth training statistics: {exc}")
+ # --- Trajectory plots ---
+ _try_plot_trajectories(config_file, config_yaml, run_dir, out_prefix, binding_df)
+
+ # --- Per-genotype theta fit plots ---
+ _try_plot_theta_fits(binding_df, pred_df, out_prefix)
+
# --- Write JSON ---
results = {
"metadata": metadata,
@@ -309,6 +766,23 @@ def summarize_fit(run_dir,
except Exception as exc:
warnings.warn(f"Could not write JSON to {json_file}: {exc}")
+ # --- Theta correlation CSVs ---
+ if train_merged is not None and len(train_merged) > 0:
+ csv_file = f"{out_prefix}_theta_corr_training.csv"
+ try:
+ train_merged.to_csv(csv_file, index=False)
+ print(f"Wrote theta training data to {csv_file}")
+ except Exception as exc:
+ warnings.warn(f"Could not write theta training CSV to {csv_file}: {exc}")
+
+ if test_merged is not None and len(test_merged) > 0:
+ csv_file = f"{out_prefix}_theta_corr_test.csv"
+ try:
+ test_merged.to_csv(csv_file, index=False)
+ print(f"Wrote theta test data to {csv_file}")
+ except Exception as exc:
+ warnings.warn(f"Could not write theta test CSV to {csv_file}: {exc}")
+
# --- Correlation plot ---
pdf_file = f"{out_prefix}_theta_corr.pdf"
try:
@@ -316,8 +790,8 @@ def summarize_fit(run_dir,
if train_merged is not None and len(train_merged) > 0:
xy_corr(
- x_values=train_merged["theta_obs"].values,
- y_values=train_merged["median"].values,
+ x_values=train_merged["ref"].values,
+ y_values=train_merged["q0.5"].values,
as_hexbin=False,
ax=axes[0],
)
@@ -329,8 +803,8 @@ def summarize_fit(run_dir,
if test_merged is not None and len(test_merged) > 0:
xy_corr(
- x_values=test_merged["theta_obs"].values,
- y_values=test_merged["median"].values,
+ x_values=test_merged["ref"].values,
+ y_values=test_merged["q0.5"].values,
as_hexbin=False,
ax=axes[1],
)
@@ -347,6 +821,23 @@ def summarize_fit(run_dir,
except Exception as exc:
warnings.warn(f"Could not generate correlation plot: {exc}")
+ # --- Growth correlation CSV ---
+ growth_copy_df = None
+ if growth_pred_file is not None:
+ growth_csv_file = f"{out_prefix}_growth_corr.csv"
+ try:
+ growth_copy_df = pd.read_csv(growth_pred_file).rename(
+ columns={"ln_cfu": "ref", "ln_cfu_std": "ref_std"}
+ )
+ growth_copy_df.to_csv(growth_csv_file, index=False)
+ print(f"Wrote growth data to {growth_csv_file}")
+ except Exception as exc:
+ warnings.warn(f"Could not write growth CSV at {growth_csv_file}: {exc}")
+ _try_calibration(
+ growth_copy_df, "ref",
+ f"{out_prefix}_growth", "growth",
+ )
+
# --- Growth correlation plot ---
growth_pdf_file = f"{out_prefix}_growth_corr.pdf"
try:
@@ -354,7 +845,7 @@ def summarize_fit(run_dir,
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
xy_corr(
x_values=growth_pred_df["ln_cfu"].values,
- y_values=growth_pred_df["median"].values,
+ y_values=growth_pred_df["q0.5"].values,
as_hexbin=False,
ax=ax,
)
@@ -373,7 +864,7 @@ def summarize_fit(run_dir,
try:
if all_losses is not None:
fig, ax = plt.subplots(figsize=(8, 4))
- ax.plot(all_epochs, all_losses, lw=1.5, color="steelblue")
+ ax.plot(all_epochs, all_losses, lw=1.5, color=DEFAULT_COLORS[2])
ax.set_xlabel("Epoch")
ax.set_ylabel("Loss")
ax.set_title("Training loss")
@@ -384,13 +875,16 @@ def summarize_fit(run_dir,
except Exception as exc:
warnings.warn(f"Could not generate loss curve plot: {exc}")
+ # --- Sim-parameter correlation CSVs ---
+ _summarize_params(run_dir, out_prefix)
+
return results
def main():
return generalized_main(
summarize_fit,
- manual_arg_types={"ground_truth_file": str},
+ manual_arg_types={"ref_theta_file": str},
)
diff --git a/src/tfscreen/tfmodel/scripts/summarize_sbc_cli.py b/src/tfscreen/tfmodel/scripts/summarize_sbc_cli.py
index 230c32fe..d105d97a 100644
--- a/src/tfscreen/tfmodel/scripts/summarize_sbc_cli.py
+++ b/src/tfscreen/tfmodel/scripts/summarize_sbc_cli.py
@@ -6,7 +6,7 @@
statistics, and writes summary outputs.
"""
-from tfscreen.tfmodel.analysis.sbc import summarize_sbc
+from tfscreen.tfmodel.analysis.error_calibration import summarize_sbc
from tfscreen.util.cli import generalized_main
diff --git a/src/tfscreen/util/__init__.py b/src/tfscreen/util/__init__.py
index f920d19e..e33eba31 100644
--- a/src/tfscreen/util/__init__.py
+++ b/src/tfscreen/util/__init__.py
@@ -1,6 +1,7 @@
from .validation import ( # noqa: F401
- check_number
+ check_number,
+ check_unknown_keys,
)
from .io import ( # noqa: F401
diff --git a/src/tfscreen/util/validation/__init__.py b/src/tfscreen/util/validation/__init__.py
index fc84a13b..033e15ca 100644
--- a/src/tfscreen/util/validation/__init__.py
+++ b/src/tfscreen/util/validation/__init__.py
@@ -1,2 +1,3 @@
from .check import check_number # noqa: F401
+from .check import check_unknown_keys # noqa: F401
diff --git a/src/tfscreen/util/validation/check.py b/src/tfscreen/util/validation/check.py
index 92e776f5..fc597640 100644
--- a/src/tfscreen/util/validation/check.py
+++ b/src/tfscreen/util/validation/check.py
@@ -4,6 +4,36 @@
# Define a generic Numeric type for hinting
_Numeric = TypeVar("_Numeric", int, float)
+
+def check_unknown_keys(
+ config: dict,
+ allowed_keys: frozenset,
+ label: str = "config",
+) -> None:
+ """
+ Raise ValueError if config contains any keys not in allowed_keys.
+
+ Parameters
+ ----------
+ config : dict
+ The configuration dictionary to validate.
+ allowed_keys : frozenset
+ The complete set of recognized top-level keys.
+ label : str
+ Human-readable name for the config type, used in the error message.
+
+ Raises
+ ------
+ ValueError
+ If any key in config is not in allowed_keys, listing all unknown keys.
+ """
+ unknown = sorted(set(config) - allowed_keys)
+ if unknown:
+ raise ValueError(
+ f"Unrecognized key(s) in {label}: {unknown}.\n"
+ f"Valid keys are: {sorted(allowed_keys)}."
+ )
+
def check_number(
value: Any,
param_name: Optional[str] = None,
diff --git a/tests/smoke-tests/test_configure_run_smoke.py b/tests/smoke-tests/test_configure_run_smoke.py
index dc5598eb..6938557c 100644
--- a/tests/smoke-tests/test_configure_run_smoke.py
+++ b/tests/smoke-tests/test_configure_run_smoke.py
@@ -4,6 +4,7 @@
import yaml
from tfscreen.tfmodel.scripts.configure_model_cli import configure_model
from tfscreen.tfmodel.scripts.fit_model_cli import fit_model
+from tfscreen.tfmodel.scripts.sample_posterior_cli import sample_posterior
from tfscreen.tfmodel.scripts.extract_params_cli import extract_params as summarize_posteriors
@pytest.mark.slow
@@ -48,12 +49,15 @@ def test_configure_run_binding_only_smoke(tmpdir):
out_prefix = os.path.join(tmpdir, "test_tfs_binding_only_out")
fit_model(config_file=config_file,
- seed=42,
- max_num_epochs=1,
- num_posterior_samples=10,
- sampling_batch_size=10,
- always_get_posterior=True,
- out_prefix=out_prefix)
+ seed=42,
+ max_num_epochs=1,
+ out_prefix=out_prefix)
+
+ sample_posterior(config_file=config_file,
+ checkpoint_file=f"{out_prefix}_checkpoint.pkl",
+ out_prefix=f"{out_prefix}_posterior",
+ num_posterior_samples=10,
+ sampling_batch_size=10)
assert os.path.exists(f"{out_prefix}_posterior.h5")
@@ -66,8 +70,8 @@ def test_configure_run_pipeline_smoke(tmpdir):
# Create mock data
growth_df = pd.DataFrame({
"replicate": ["R1", "R1"],
- "condition_pre": ["C1", "C1"],
- "condition_sel": ["C2", "C2"],
+ "condition_pre": ["C1-", "C1-"],
+ "condition_sel": ["C2+", "C2+"],
"genotype": ["A123B", "C456D"],
"t": [0, 10],
"t_sel": [0, 10],
@@ -114,17 +118,18 @@ def test_configure_run_pipeline_smoke(tmpdir):
assert "priors" not in config
assert "init_params" not in config
- # Run analysis (smoke test)
- # We use max_num_epochs=1 to make it fast
- # We use always_get_posterior=True to ensure the posterior file is written
+ # Run analysis (smoke test); use max_num_epochs=1 to make it fast
out_prefix = os.path.join(tmpdir, "test_tfs_out")
fit_model(config_file=config_file,
- seed=42,
- max_num_epochs=1,
- num_posterior_samples=10,
- sampling_batch_size=10,
- always_get_posterior=True,
- out_prefix=out_prefix)
+ seed=42,
+ max_num_epochs=1,
+ out_prefix=out_prefix)
+
+ sample_posterior(config_file=config_file,
+ checkpoint_file=f"{out_prefix}_checkpoint.pkl",
+ out_prefix=f"{out_prefix}_posterior",
+ num_posterior_samples=10,
+ sampling_batch_size=10)
assert os.path.exists(os.path.join(tmpdir, "test_tfs_out_posterior.h5"))
@@ -150,8 +155,8 @@ def test_configure_run_binding_weight_smoke(tmpdir):
# Growth data: 20 rows | Binding data: 4 rows → expected auto-weight = 5.0
growth_df = pd.DataFrame({
"replicate": ["R1"] * 20,
- "condition_pre": ["C1"] * 20,
- "condition_sel": ["C2"] * 20,
+ "condition_pre": ["C1-"] * 20,
+ "condition_sel": ["C2+"] * 20,
"genotype": ["A123B"] * 10 + ["C456D"] * 10,
"t_sel": [10.0] * 20,
"t_pre": [12.0] * 20,
@@ -205,9 +210,12 @@ def test_configure_run_binding_weight_smoke(tmpdir):
fit_model(config_file=f"{out_auto}_config.yaml",
seed=42,
max_num_epochs=1,
- num_posterior_samples=10,
- sampling_batch_size=10,
- always_get_posterior=True,
out_prefix=out_fit)
+ sample_posterior(config_file=f"{out_auto}_config.yaml",
+ checkpoint_file=f"{out_fit}_checkpoint.pkl",
+ out_prefix=f"{out_fit}_posterior",
+ num_posterior_samples=10,
+ sampling_batch_size=10)
+
assert os.path.exists(f"{out_fit}_posterior.h5")
diff --git a/tests/smoke-tests/test_prediction_smoke.py b/tests/smoke-tests/test_prediction_smoke.py
index b3235a12..2000d69c 100644
--- a/tests/smoke-tests/test_prediction_smoke.py
+++ b/tests/smoke-tests/test_prediction_smoke.py
@@ -68,9 +68,9 @@ def test_prediction_smoke(growth_smoke_csv,
assert set(pred_df["titrant_conc"]) == set(new_titrant_conc)
# Check for quantile columns (defaults from load_posteriors)
- assert "median" in pred_df.columns
- assert "lower_95" in pred_df.columns
- assert "upper_95" in pred_df.columns
+ assert "q0.5" in pred_df.columns
+ assert "q0.025" in pred_df.columns
+ assert "q0.975" in pred_df.columns
# 3. Test expansion restriction for plated dimensions
# 'logit_norm' plates on titrant_conc, so expanding it should fail.
diff --git a/tests/smoke-tests/test_run_growth_analysis_smoke.py b/tests/smoke-tests/test_run_growth_analysis_smoke.py
index 525d0125..e1d69499 100644
--- a/tests/smoke-tests/test_run_growth_analysis_smoke.py
+++ b/tests/smoke-tests/test_run_growth_analysis_smoke.py
@@ -7,11 +7,11 @@
from tfscreen.tfmodel.scripts.sample_posterior_cli import sample_posterior
-def _build_config(gm, tmpdir, growth_smoke_csv, binding_smoke_csv):
- """Write a config file for gm and return its path."""
+def _build_config(orchestrator, tmpdir, growth_smoke_csv, binding_smoke_csv):
+ """Write a config file for orchestrator and return its path."""
config_root = os.path.join(tmpdir, "smoke")
write_configuration(
- gm=gm,
+ orchestrator=orchestrator,
out_prefix=config_root,
growth_df_path=str(growth_smoke_csv),
binding_df_path=str(binding_smoke_csv),
@@ -26,7 +26,7 @@ def test_run_growth_analysis_smoke(growth_smoke_csv, binding_smoke_csv, tmpdir):
out_prefix = os.path.join(tmpdir, "run_smoke")
# 1. Initialize a model and write its configuration
- gm = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
growth_df=growth_smoke_csv,
binding_df=binding_smoke_csv,
condition_growth="linear",
@@ -37,7 +37,7 @@ def test_run_growth_analysis_smoke(growth_smoke_csv, binding_smoke_csv, tmpdir):
config_root = os.path.join(tmpdir, "smoke")
write_configuration(
- gm=gm,
+ orchestrator=orchestrator,
out_prefix=config_root,
growth_df_path=str(growth_smoke_csv),
binding_df_path=str(binding_smoke_csv)
@@ -54,16 +54,22 @@ def test_run_growth_analysis_smoke(growth_smoke_csv, binding_smoke_csv, tmpdir):
analysis_method="svi",
out_prefix=out_prefix,
max_num_epochs=1,
+ forward_batch_size=10,
+ )
+
+ sample_posterior(
+ config_file=config_file,
+ checkpoint_file=f"{out_prefix}_checkpoint.pkl",
+ out_prefix=f"{out_prefix}_posterior",
num_posterior_samples=5,
sampling_batch_size=5,
forward_batch_size=10,
- always_get_posterior=True
)
-
+
# 3. Verify outputs
posterior_file = f"{out_prefix}_posterior.h5"
assert os.path.exists(posterior_file)
-
+
# Load and check the saved file
with h5py.File(posterior_file, 'r') as data:
assert "growth_pred" in data
@@ -79,7 +85,7 @@ def test_run_growth_analysis_nuts_smoke(growth_smoke_csv, binding_smoke_csv, tmp
- a posterior HDF5 is written with the expected keys
- a checkpoint pkl is written containing 'mcmc_samples'
"""
- gm = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
growth_df=growth_smoke_csv,
binding_df=binding_smoke_csv,
condition_growth="linear",
@@ -87,7 +93,7 @@ def test_run_growth_analysis_nuts_smoke(growth_smoke_csv, binding_smoke_csv, tmp
theta="hill_geno",
batch_size=None,
)
- config_file = _build_config(gm, tmpdir, growth_smoke_csv, binding_smoke_csv)
+ config_file = _build_config(orchestrator, tmpdir, growth_smoke_csv, binding_smoke_csv)
assert os.path.exists(config_file)
out_prefix = os.path.join(tmpdir, "nuts_smoke")
@@ -134,7 +140,7 @@ def test_run_growth_analysis_nuts_posterior_smoke(growth_smoke_csv, binding_smok
Runs a tiny NUTS chain, then re-generates posteriors from the checkpoint
and verifies that the output HDF5 is recreated correctly.
"""
- gm = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
growth_df=growth_smoke_csv,
binding_df=binding_smoke_csv,
condition_growth="linear",
@@ -142,7 +148,7 @@ def test_run_growth_analysis_nuts_posterior_smoke(growth_smoke_csv, binding_smok
theta="hill_geno",
batch_size=None,
)
- config_file = _build_config(gm, tmpdir, growth_smoke_csv, binding_smoke_csv)
+ config_file = _build_config(orchestrator, tmpdir, growth_smoke_csv, binding_smoke_csv)
# Step 1: run NUTS to produce a checkpoint
nuts_root = os.path.join(tmpdir, "nuts_for_post")
diff --git a/tests/tfscreen/tfmodel/components/dk_geno/__init__.py b/tests/tfscreen/analysis/cat_response/scripts/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/dk_geno/__init__.py
rename to tests/tfscreen/analysis/cat_response/scripts/__init__.py
diff --git a/tests/tfscreen/analysis/cat_response/test_cat_response_cli.py b/tests/tfscreen/analysis/cat_response/scripts/test_cat_response_cli.py
similarity index 70%
rename from tests/tfscreen/analysis/cat_response/test_cat_response_cli.py
rename to tests/tfscreen/analysis/cat_response/scripts/test_cat_response_cli.py
index 2081f87a..d47051a8 100644
--- a/tests/tfscreen/analysis/cat_response/test_cat_response_cli.py
+++ b/tests/tfscreen/analysis/cat_response/scripts/test_cat_response_cli.py
@@ -5,7 +5,7 @@
import pandas as pd
from unittest.mock import patch
-from tfscreen.analysis.cat_response.cat_response_cli import cat_response
+from tfscreen.analysis.cat_response.scripts.cat_response_cli import cat_response
# Run ProcessPoolExecutor synchronously so cat_fit mocks work in-process.
@@ -22,7 +22,7 @@ def __exit__(self, *a): pass
def submit(self, fn, *args): return _SyncFuture(fn(*args))
_SYNC_EXECUTOR = patch(
- "tfscreen.analysis.cat_response.cat_response_cli.ProcessPoolExecutor",
+ "tfscreen.analysis.cat_response.scripts.cat_response_cli.ProcessPoolExecutor",
return_value=_SyncExecutor(),
)
@@ -36,15 +36,15 @@ def _make_theta_df(center_col):
"titrant_name": ["IPTG", "IPTG"],
"titrant_conc": [0.0, 1.0],
center_col: [0.3, 0.7],
- "upper_std": [0.4, 0.8],
- "lower_std": [0.2, 0.6],
+ "q0.841": [0.4, 0.8],
+ "q0.159": [0.2, 0.6],
})
class TestThetaColAutoDetect:
- def test_autodetects_median(self, tmp_path):
- """Values from median column are passed to the fitter."""
+ def test_autodetects_q0_5(self, tmp_path):
+ """Values from q0.5 column are passed to the fitter."""
f = str(tmp_path / "theta.csv")
captured = {}
@@ -52,10 +52,10 @@ def fake_fit(x, y, y_std, models_to_run):
captured["y"] = list(y)
return (_FLAT_RESULT, None)
- df = _make_theta_df("median")
- with patch("tfscreen.analysis.cat_response.cat_response_cli.pd.read_csv",
+ df = _make_theta_df("q0.5")
+ with patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.pd.read_csv",
return_value=df), \
- patch("tfscreen.analysis.cat_response.cat_response_cli.cat_fit",
+ patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.cat_fit",
side_effect=fake_fit), \
_SYNC_EXECUTOR:
cat_response(f, out_prefix=str(tmp_path / "out"))
@@ -72,9 +72,9 @@ def fake_fit(x, y, y_std, models_to_run):
return (_FLAT_RESULT, None)
df = _make_theta_df("point_est")
- with patch("tfscreen.analysis.cat_response.cat_response_cli.pd.read_csv",
+ with patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.pd.read_csv",
return_value=df), \
- patch("tfscreen.analysis.cat_response.cat_response_cli.cat_fit",
+ patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.cat_fit",
side_effect=fake_fit), \
_SYNC_EXECUTOR:
cat_response(f, out_prefix=str(tmp_path / "out"))
@@ -90,19 +90,19 @@ def fake_fit(x, y, y_std, models_to_run):
captured["y"] = list(y)
return (_FLAT_RESULT, None)
- df = _make_theta_df("median").copy()
+ df = _make_theta_df("q0.5").copy()
df["my_col"] = [0.11, 0.22]
- with patch("tfscreen.analysis.cat_response.cat_response_cli.pd.read_csv",
+ with patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.pd.read_csv",
return_value=df), \
- patch("tfscreen.analysis.cat_response.cat_response_cli.cat_fit",
+ patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.cat_fit",
side_effect=fake_fit), \
_SYNC_EXECUTOR:
cat_response(f, theta_col="my_col", out_prefix=str(tmp_path / "out"))
assert captured["y"] == pytest.approx([0.11, 0.22])
- def test_median_preferred_over_point_est(self, tmp_path):
- """When both median and point_est are present, median wins."""
+ def test_q0_5_preferred_over_point_est(self, tmp_path):
+ """When both q0.5 and point_est are present, q0.5 wins."""
f = str(tmp_path / "theta.csv")
captured = {}
@@ -110,11 +110,11 @@ def fake_fit(x, y, y_std, models_to_run):
captured["y"] = list(y)
return (_FLAT_RESULT, None)
- df = _make_theta_df("median").copy()
+ df = _make_theta_df("q0.5").copy()
df["point_est"] = [0.9, 0.8]
- with patch("tfscreen.analysis.cat_response.cat_response_cli.pd.read_csv",
+ with patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.pd.read_csv",
return_value=df), \
- patch("tfscreen.analysis.cat_response.cat_response_cli.cat_fit",
+ patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.cat_fit",
side_effect=fake_fit), \
_SYNC_EXECUTOR:
cat_response(f, out_prefix=str(tmp_path / "out"))
@@ -122,16 +122,16 @@ def fake_fit(x, y, y_std, models_to_run):
assert captured["y"] == pytest.approx([0.3, 0.7])
def test_raises_when_no_theta_col(self, tmp_path):
- """ValueError when neither median nor point_est is present."""
+ """ValueError when neither q0.5 nor point_est is present."""
f = str(tmp_path / "theta.csv")
df = pd.DataFrame({
"genotype": ["wt"],
"titrant_name": ["IPTG"],
"titrant_conc": [0.0],
- "upper_std": [0.5],
- "lower_std": [0.3],
+ "q0.841": [0.5],
+ "q0.159": [0.3],
})
- with patch("tfscreen.analysis.cat_response.cat_response_cli.pd.read_csv",
+ with patch("tfscreen.analysis.cat_response.scripts.cat_response_cli.pd.read_csv",
return_value=df):
with pytest.raises(ValueError, match="No theta column found"):
cat_response(f, out_prefix=str(tmp_path / "out"))
diff --git a/tests/tfscreen/analysis/cat_response/test_fit_response_cli.py b/tests/tfscreen/analysis/cat_response/test_fit_response_cli.py
deleted file mode 100644
index 6fc0fc26..00000000
--- a/tests/tfscreen/analysis/cat_response/test_fit_response_cli.py
+++ /dev/null
@@ -1,116 +0,0 @@
-"""
-Tests for fit_response_cli.py — theta_col auto-detection and error handling.
-"""
-import pytest
-import pandas as pd
-from unittest.mock import patch
-
-from tfscreen.analysis.cat_response.fit_response_cli import fit_response
-
-
-# Run ProcessPoolExecutor synchronously so cat_fit mocks work in-process.
-from concurrent.futures import Future as _Future
-
-class _SyncFuture(_Future):
- def __init__(self, result_val):
- super().__init__()
- self.set_result(result_val)
-
-class _SyncExecutor:
- def __enter__(self): return self
- def __exit__(self, *a): pass
- def submit(self, fn, *args): return _SyncFuture(fn(*args))
-
-_SYNC_EXECUTOR = patch(
- "tfscreen.analysis.cat_response.fit_response_cli.ProcessPoolExecutor",
- return_value=_SyncExecutor(),
-)
-
-_FLAT_RESULT = {"genotype": "wt", "titrant_name": "IPTG",
- "best_model": "flat", "status": "success"}
-
-
-def _make_df(center_col):
- return pd.DataFrame({
- "genotype": ["wt", "wt"],
- "titrant_name": ["IPTG", "IPTG"],
- "titrant_conc": [0.0, 1.0],
- center_col: [0.3, 0.7],
- "upper_std": [0.4, 0.8],
- "lower_std": [0.2, 0.6],
- })
-
-
-class TestFitResponseThetaColAutoDetect:
-
- def test_autodetects_median(self):
- """Values from median column are passed to the fitter."""
- captured = {}
-
- def fake_fit(x, y, y_std, models_to_run):
- captured["y"] = list(y)
- return (_FLAT_RESULT, None)
-
- with patch("tfscreen.analysis.cat_response.fit_response_cli.cat_fit",
- side_effect=fake_fit), _SYNC_EXECUTOR:
- fit_response(_make_df("median"), models_to_run=["flat"])
-
- assert captured["y"] == pytest.approx([0.3, 0.7])
-
- def test_autodetects_point_est(self):
- """Values from point_est column are passed to the fitter when median is absent."""
- captured = {}
-
- def fake_fit(x, y, y_std, models_to_run):
- captured["y"] = list(y)
- return (_FLAT_RESULT, None)
-
- with patch("tfscreen.analysis.cat_response.fit_response_cli.cat_fit",
- side_effect=fake_fit), _SYNC_EXECUTOR:
- fit_response(_make_df("point_est"), models_to_run=["flat"])
-
- assert captured["y"] == pytest.approx([0.3, 0.7])
-
- def test_explicit_theta_col_overrides(self):
- """Explicit theta_col is used even when median is also present."""
- captured = {}
-
- def fake_fit(x, y, y_std, models_to_run):
- captured["y"] = list(y)
- return (_FLAT_RESULT, None)
-
- df = _make_df("median").copy()
- df["my_col"] = [0.11, 0.22]
- with patch("tfscreen.analysis.cat_response.fit_response_cli.cat_fit",
- side_effect=fake_fit), _SYNC_EXECUTOR:
- fit_response(df, theta_col="my_col", models_to_run=["flat"])
-
- assert captured["y"] == pytest.approx([0.11, 0.22])
-
- def test_median_preferred_over_point_est(self):
- """When both median and point_est are present, median wins."""
- captured = {}
-
- def fake_fit(x, y, y_std, models_to_run):
- captured["y"] = list(y)
- return (_FLAT_RESULT, None)
-
- df = _make_df("median").copy()
- df["point_est"] = [0.9, 0.8]
- with patch("tfscreen.analysis.cat_response.fit_response_cli.cat_fit",
- side_effect=fake_fit), _SYNC_EXECUTOR:
- fit_response(df, models_to_run=["flat"])
-
- assert captured["y"] == pytest.approx([0.3, 0.7])
-
- def test_raises_when_no_theta_col(self):
- """ValueError when neither median nor point_est is present."""
- df = pd.DataFrame({
- "genotype": ["wt"],
- "titrant_name": ["IPTG"],
- "titrant_conc": [0.0],
- "upper_std": [0.5],
- "lower_std": [0.3],
- })
- with pytest.raises(ValueError, match="No theta column found"):
- fit_response(df, models_to_run=["flat"])
diff --git a/tests/tfscreen/mle/test_stats_test_suite.py b/tests/tfscreen/analysis/test_stats_test_suite.py
similarity index 95%
rename from tests/tfscreen/mle/test_stats_test_suite.py
rename to tests/tfscreen/analysis/test_stats_test_suite.py
index 7bfae6df..1f294ebd 100644
--- a/tests/tfscreen/mle/test_stats_test_suite.py
+++ b/tests/tfscreen/analysis/test_stats_test_suite.py
@@ -1,7 +1,7 @@
import pytest
import numpy as np
-from tfscreen.mle.stats_test_suite import stats_test_suite
+from tfscreen.analysis.stats_test_suite import stats_test_suite
def test_stats_test_suite_perfect_fit():
param_real = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
@@ -86,7 +86,7 @@ def test_stats_test_suite_no_std():
import warnings
def test_het_breuschpagan_linalg_error(mocker):
# Mock het_breuschpagan to raise LinAlgError
- mocker.patch("tfscreen.mle.stats_test_suite.het_breuschpagan", side_effect=np.linalg.LinAlgError)
+ mocker.patch("tfscreen.analysis.stats_test_suite.het_breuschpagan", side_effect=np.linalg.LinAlgError)
param_real = np.array([1.0, 2.0, 3.0])
param_est = np.array([1.1, 2.2, 3.3])
diff --git a/tests/tfscreen/genetics/test_library_manager.py b/tests/tfscreen/genetics/test_library_manager.py
index 957ff412..333c3098 100644
--- a/tests/tfscreen/genetics/test_library_manager.py
+++ b/tests/tfscreen/genetics/test_library_manager.py
@@ -183,6 +183,18 @@ def test_parse_success(valid_config):
assert lm.libraries_seen == {"2"}
assert lm.degen_seq == "gattacnnncgattaca"
+def test_parse_whitespace_stripped(valid_config):
+ """
+ Whitespace (spaces, newlines) in wt_seq, degen_sites, and sub_libraries
+ is removed before validation, allowing YAML literal-block multiline values.
+ """
+ valid_config["wt_seq"] = "gattac\nagtcgattaca"
+ valid_config["degen_sites"] = "......\nnnn........"
+ valid_config["sub_libraries"] = "......\n222........"
+ lm = LibraryManager(valid_config)
+ assert lm.expected_length == 17
+ assert lm.wt_seq == "gattacagtcgattaca"
+
def test_parse_missing_key(valid_config):
"""
Tests that a ValueError is raised if a required key is missing.
@@ -676,6 +688,7 @@ def lm_for_spiked_seqs() -> LibraryManager:
# Set attributes required by the method under test
lm.wt_seq = "gattaca"
lm.standard_bases = set('acgt')
+ lm.standard_plus_dot = set('acgt.')
return lm
def test_get_spiked_seqs_success(lm_for_spiked_seqs):
@@ -704,15 +717,42 @@ def test_get_spiked_seqs_success(lm_for_spiked_seqs):
assert returned_seqs == valid_seqs
assert returned_aa == ["", "Y3K"]
+def test_get_spiked_seqs_dot_expansion(lm_for_spiked_seqs):
+ """
+ Dots in spiked sequences are replaced by the corresponding wt_seq base.
+ """
+ # wt_seq = "gattaca"; "gat.acg" -> dot at pos 3 fills in wt 't'; 'g' at
+ # pos 6 is an explicit mutation (wt is 'a') -> "gattacg"
+ dot_seq = ["gat.acg"]
+
+ with patch.object(lm_for_spiked_seqs, '_convert_to_aa') as mock_convert:
+ mock_convert.return_value = ["A6G"]
+ returned_seqs, _ = lm_for_spiked_seqs._get_spiked_seqs(dot_seq)
+
+ assert returned_seqs == ["gattacg"]
+
+def test_get_spiked_seqs_whitespace_stripped(lm_for_spiked_seqs):
+ """
+ Whitespace (spaces, newlines) in spiked sequences is stripped before
+ validation, allowing YAML literal-block multiline entries.
+ """
+ # wt_seq = "gattaca"; "gat\ntaca" has an embedded newline that should be
+ # removed, yielding "gattaca" (7 chars, matches wt length)
+ with patch.object(lm_for_spiked_seqs, '_convert_to_aa') as mock_convert:
+ mock_convert.return_value = [""]
+ returned_seqs, _ = lm_for_spiked_seqs._get_spiked_seqs(["gat\ntaca"])
+
+ assert returned_seqs == ["gattaca"]
+
def test_get_spiked_seqs_invalid_character(lm_for_spiked_seqs):
"""
Tests that a ValueError is raised for a sequence with non-standard bases.
"""
invalid_seqs = ["gattaca", "gattacx"] # Contains 'x'
-
+
with pytest.raises(ValueError) as excinfo:
lm_for_spiked_seqs._get_spiked_seqs(invalid_seqs)
-
+
assert "Characters not recognized" in str(excinfo.value)
assert "spiked seq" in str(excinfo.value)
diff --git a/tests/tfscreen/plot/heatmap/test_heatmap_core.py b/tests/tfscreen/plot/heatmap/test_heatmap_core.py
index a7125ba7..bb350827 100644
--- a/tests/tfscreen/plot/heatmap/test_heatmap_core.py
+++ b/tests/tfscreen/plot/heatmap/test_heatmap_core.py
@@ -641,10 +641,24 @@ def mock_helpers(self, mocker) -> dict:
"format_axis": mocker.patch("tfscreen.plot.heatmap.heatmap_core._format_axis"),
"colorbar": mocker.patch("matplotlib.figure.Figure.colorbar")
}
- # The collection needs to be a mock that can be added to an ax
+ # The collection needs to be a mock that can be added to an ax.
+ # get_transform().contains_branch_separately() must return a 2-tuple so
+ # matplotlib's _update_collection_limits can unpack (x_is_data, y_is_data).
+ # The method was misspelled in matplotlib <3.10 as "seperately"; cover both.
mock_p_collection = MagicMock(spec=PatchCollection)
- mock_p_collection.get_transform().contains_branch_seperately.return_value = (True, True)
- mock_p_collection.get_offset_transform().contains_branch_seperately.return_value = (True, True)
+ mock_transform = MagicMock()
+ mock_transform.contains_branch_separately.return_value = (True, True)
+ mock_transform.contains_branch_seperately.return_value = (True, True)
+ mock_p_collection.get_transform.return_value = mock_transform
+ mock_p_collection.get_offset_transform.return_value = mock_transform
+ # _sticky_edges is an instance attribute (set in Artist.__init__) so it
+ # is absent from a class-spec'd mock; set it explicitly with the shape
+ # matplotlib expects: an object with .x and .y list attributes.
+ mock_sticky = MagicMock()
+ mock_sticky.x = []
+ mock_sticky.y = []
+ mock_p_collection._sticky_edges = mock_sticky
+ mock_p_collection.sticky_edges = mock_sticky
mocks["build_heatmap_collection"].return_value = mock_p_collection
return mocks
diff --git a/tests/tfscreen/plot/test_geno_trajectory.py b/tests/tfscreen/plot/test_geno_trajectory.py
new file mode 100644
index 00000000..1cee409e
--- /dev/null
+++ b/tests/tfscreen/plot/test_geno_trajectory.py
@@ -0,0 +1,551 @@
+"""
+Tests for tfscreen.plot.geno_trajectory.
+
+plot_geno_trajectory is a pure plotting function — all tests pass synthetic
+DataFrames and inspect the returned Figure without requiring JAX or a real
+ModelOrchestrator.
+
+predict_and_plot_geno_trajectory requires a live orchestrator and is covered
+by smoke tests; only its .npz loading path is tested here via a mock.
+"""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from tfscreen.plot.geno_trajectory import plot_geno_trajectory
+
+
+# ---------------------------------------------------------------------------
+# Synthetic DataFrame helpers
+# ---------------------------------------------------------------------------
+
+def _make_pred_df(
+ genotypes=("geno_A", "geno_B"),
+ replicates=("rep1", "rep2"),
+ conditions=(("pre", "sel", "drug", 0.0), ("pre", "sel", "drug", 1.0)),
+ t_values=(-5.0, 0.0, 5.0, 10.0, 15.0),
+ include_ci=True,
+ include_obs=True,
+):
+ """
+ Build a minimal synthetic pred_df.
+
+ Observed data (non-NaN ln_cfu) is placed at t_values that are >= 0.
+ Fine-grid rows have NaN ln_cfu. Rows at negative t carry ln_cfu0 quantiles
+ (simulating the pre-selection anchor).
+ """
+ rng = np.random.default_rng(42)
+ rows = []
+ for cp, cs, tn, tc in conditions:
+ for geno in genotypes:
+ for rep in replicates:
+ base = rng.uniform(8.0, 12.0)
+ slope = rng.uniform(0.1, 0.3)
+ for t in t_values:
+ med = base + slope * t
+ row = {
+ "replicate": rep,
+ "condition_pre": cp,
+ "condition_sel": cs,
+ "titrant_name": tn,
+ "titrant_conc": tc,
+ "genotype": geno,
+ "t_sel": t,
+ "ln_cfu": (base + slope * t) if (include_obs and t >= 0) else np.nan,
+ "ln_cfu_std": 0.2 if (include_obs and t >= 0) else np.nan,
+ "q0.5": med,
+ }
+ if include_ci:
+ row["q0.05"] = med - 0.5
+ row["q0.95"] = med + 0.5
+ rows.append(row)
+ return pd.DataFrame(rows)
+
+
+# ---------------------------------------------------------------------------
+# Return type
+# ---------------------------------------------------------------------------
+
+class TestReturnType:
+ def test_returns_figure(self):
+ df = _make_pred_df()
+ fig = plot_geno_trajectory(df)
+ assert isinstance(fig, plt.Figure)
+ plt.close(fig)
+
+ def test_figure_has_axes(self):
+ df = _make_pred_df()
+ fig = plot_geno_trajectory(df)
+ assert len(fig.axes) > 0
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Subplot count
+# ---------------------------------------------------------------------------
+
+class TestSubplotCount:
+ def test_one_subplot_per_condition_combo(self):
+ n_conds = 3
+ conditions = tuple(
+ ("pre", "sel", "drug", float(i)) for i in range(n_conds)
+ )
+ df = _make_pred_df(conditions=conditions)
+ fig = plot_geno_trajectory(df)
+ visible = [ax for ax in fig.axes if ax.get_visible()]
+ assert len(visible) == n_conds
+ plt.close(fig)
+
+ def test_single_condition_one_subplot(self):
+ df = _make_pred_df(conditions=(("pre", "sel", "drug", 0.0),))
+ fig = plot_geno_trajectory(df)
+ visible = [ax for ax in fig.axes if ax.get_visible()]
+ assert len(visible) == 1
+ plt.close(fig)
+
+ def test_four_conditions_two_rows(self):
+ conditions = tuple(
+ ("pre", "sel", "drug", float(i)) for i in range(4)
+ )
+ df = _make_pred_df(conditions=conditions)
+ fig = plot_geno_trajectory(df)
+ n_rows = fig.axes[0].get_subplotspec().get_gridspec().nrows
+ assert n_rows == 2
+ plt.close(fig)
+
+ def test_unused_cells_hidden(self):
+ # 4 conditions → 2×3 grid → 2 hidden cells
+ conditions = tuple(
+ ("pre", "sel", "drug", float(i)) for i in range(4)
+ )
+ df = _make_pred_df(conditions=conditions)
+ fig = plot_geno_trajectory(df)
+ hidden = [ax for ax in fig.axes if not ax.get_visible()]
+ assert len(hidden) == 2
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Observed data and predicted lines
+# ---------------------------------------------------------------------------
+
+class TestObservedAndPredicted:
+ def test_no_crash_all_nan_ln_cfu(self):
+ df = _make_pred_df()
+ df["ln_cfu"] = np.nan
+ df["ln_cfu_std"] = np.nan
+ fig = plot_geno_trajectory(df)
+ assert isinstance(fig, plt.Figure)
+ plt.close(fig)
+
+ def test_no_crash_all_nan_median(self):
+ df = _make_pred_df()
+ df["q0.5"] = np.nan
+ with pytest.warns(UserWarning, match="No artists with labels"):
+ fig = plot_geno_trajectory(df)
+ assert isinstance(fig, plt.Figure)
+ plt.close(fig)
+
+ def test_axvline_at_zero_present(self):
+ df = _make_pred_df(conditions=(("pre", "sel", "drug", 0.0),))
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ # The axvline draws a vertical line; its x-data is (0, 0).
+ vlines = [l for l in ax.lines if list(l.get_xdata()) == [0.0, 0.0]]
+ assert len(vlines) >= 1
+ plt.close(fig)
+
+ def test_predicted_lines_drawn(self):
+ """Each (genotype, replicate) pair should produce at least one line."""
+ n_geno = 2
+ n_rep = 2
+ genotypes = tuple(f"g{i}" for i in range(n_geno))
+ replicates = tuple(f"r{i}" for i in range(n_rep))
+ df = _make_pred_df(
+ genotypes=genotypes,
+ replicates=replicates,
+ conditions=(("pre", "sel", "drug", 0.0),),
+ include_ci=False,
+ )
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ # axvline + one line per (geno, rep) pair
+ assert len(ax.lines) >= n_geno * n_rep + 1
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Credible interval
+# ---------------------------------------------------------------------------
+
+class TestCredibleInterval:
+ def test_ci_band_present_when_columns_exist(self):
+ df = _make_pred_df(include_ci=True)
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ # fill_between produces PolyCollection objects
+ assert len(ax.collections) > 0
+ plt.close(fig)
+
+ def test_no_ci_band_without_columns(self):
+ df = _make_pred_df(include_ci=False)
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ # fill_between produces PolyCollection; errorbar may add LineCollections.
+ # Only PolyCollections indicate a CI band.
+ from matplotlib.collections import PolyCollection
+ poly_cols = [c for c in ax.collections if isinstance(c, PolyCollection)]
+ assert len(poly_cols) == 0
+ plt.close(fig)
+
+ def test_ci_spans_negative_t(self):
+ """CI band should cover pre-selection (t < 0) anchor rows."""
+ df = _make_pred_df(
+ conditions=(("pre", "sel", "drug", 0.0),),
+ genotypes=("g0",),
+ replicates=("r0",),
+ t_values=(-5.0, 0.0, 5.0),
+ include_ci=True,
+ )
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ polys = ax.collections
+ assert len(polys) > 0
+ # The polygon should include x-values less than zero
+ x_vals = np.concatenate([p.get_paths()[0].vertices[:, 0] for p in polys])
+ assert np.any(x_vals < 0)
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Legend
+# ---------------------------------------------------------------------------
+
+class TestLegend:
+ def test_legend_contains_genotype_replicate(self):
+ df = _make_pred_df(
+ genotypes=("WT", "mut1"),
+ replicates=("rep1",),
+ conditions=(("pre", "sel", "drug", 0.0),),
+ )
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ legend_texts = [t.get_text() for t in ax.get_legend().get_texts()]
+ assert any("WT" in t for t in legend_texts)
+ assert any("mut1" in t for t in legend_texts)
+ plt.close(fig)
+
+ def test_legend_includes_replicate_label(self):
+ df = _make_pred_df(
+ genotypes=("WT",),
+ replicates=("rep1", "rep2"),
+ conditions=(("pre", "sel", "drug", 0.0),),
+ )
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ legend_texts = [t.get_text() for t in ax.get_legend().get_texts()]
+ assert any("rep1" in t for t in legend_texts)
+ assert any("rep2" in t for t in legend_texts)
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Styling overrides
+# ---------------------------------------------------------------------------
+
+class TestStylingOverrides:
+ def test_custom_figsize(self):
+ df = _make_pred_df(conditions=(("pre", "sel", "drug", 0.0),))
+ fig = plot_geno_trajectory(df, figsize=(10, 8))
+ w, h = fig.get_size_inches()
+ assert pytest.approx(w) == 10
+ assert pytest.approx(h) == 8
+ plt.close(fig)
+
+ def test_default_figsize_scales_with_columns(self):
+ # 3 conditions → n_cols=3 → width=15
+ conditions = tuple(
+ ("pre", "sel", "drug", float(i)) for i in range(3)
+ )
+ df = _make_pred_df(conditions=conditions)
+ fig = plot_geno_trajectory(df)
+ w, _ = fig.get_size_inches()
+ assert pytest.approx(w) == 15
+ plt.close(fig)
+
+ def test_custom_colors_accepted(self):
+ df = _make_pred_df(genotypes=("g0", "g1"))
+ fig = plot_geno_trajectory(df, colors=["#ff0000", "#0000ff"])
+ assert isinstance(fig, plt.Figure)
+ plt.close(fig)
+
+ def test_custom_markers_accepted(self):
+ df = _make_pred_df(replicates=("r0", "r1"))
+ fig = plot_geno_trajectory(df, markers=["o", "s"])
+ assert isinstance(fig, plt.Figure)
+ plt.close(fig)
+
+ def test_genotype_replicate_pairs_get_distinct_colors(self):
+ """Each (genotype, replicate) pair should get a distinct color."""
+ df = _make_pred_df(
+ genotypes=("g0", "g1"),
+ replicates=("r0", "r1"),
+ conditions=(("pre", "sel", "drug", 0.0),),
+ include_ci=False,
+ )
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ # Predicted lines are drawn with lw=1.8; exclude axvline
+ series_lines = [l for l in ax.lines if l.get_linewidth() == pytest.approx(1.8)]
+ colors_used = {l.get_color() for l in series_lines}
+ # 2 genotypes × 2 replicates = 4 pairs → 4 distinct colors
+ assert len(colors_used) == 4
+ plt.close(fig)
+
+ def test_same_genotype_different_replicates_different_colors(self):
+ """Two replicates of the same genotype must get different colors."""
+ df = _make_pred_df(
+ genotypes=("g0",),
+ replicates=("r0", "r1"),
+ conditions=(("pre", "sel", "drug", 0.0),),
+ include_ci=False,
+ )
+ fig = plot_geno_trajectory(df)
+ ax = next(ax for ax in fig.axes if ax.get_visible())
+ series_lines = [l for l in ax.lines if l.get_linewidth() == pytest.approx(1.8)]
+ colors_used = {l.get_color() for l in series_lines}
+ assert len(colors_used) == 2
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# predict_and_plot_geno_trajectory — .npz loading path (mocked)
+# ---------------------------------------------------------------------------
+
+class TestPredictAndPlotNpzLoading:
+ def test_npz_path_converted_to_dict(self, tmp_path):
+ """Passing a .npz path loads the file before calling predict()."""
+ from unittest.mock import MagicMock, patch
+ import matplotlib
+ matplotlib.use("Agg")
+
+ npz_path = str(tmp_path / "params.npz")
+ np.savez(npz_path, x_auto_loc=np.array(1.0))
+
+ # Build a minimal fake orchestrator
+ gd = MagicMock()
+ gd.good_mask = np.ones((1, 3, 1, 1, 1, 1, 1), dtype=bool)
+ gd.t_sel = np.broadcast_to(
+ np.array([0.0, 5.0, 10.0])[None, :, None, None, None, None, None],
+ (1, 3, 1, 1, 1, 1, 1),
+ ).copy()
+ data = MagicMock()
+ data.growth = gd
+
+ orchestrator = MagicMock()
+ orchestrator.data = data
+ orchestrator.presplit_df = None
+
+ # Synthetic fine_df / ln_cfu0_df to return from mocked predict()
+ fine_df = pd.DataFrame({
+ "replicate": ["r0"] * 5,
+ "condition_pre": ["pre"] * 5,
+ "condition_sel": ["sel"] * 5,
+ "titrant_name": ["drug"] * 5,
+ "titrant_conc": [0.0] * 5,
+ "genotype": ["g0"] * 5,
+ "t_sel": np.linspace(0, 10, 5),
+ "ln_cfu": [np.nan] * 5,
+ "ln_cfu_std": [np.nan] * 5,
+ "q0.05": np.linspace(8, 12, 5) - 0.5,
+ "q0.5": np.linspace(8, 12, 5),
+ "q0.95": np.linspace(8, 12, 5) + 0.5,
+ })
+ orchestrator.growth_df = pd.DataFrame({
+ "replicate": ["r0"],
+ "condition_pre": ["pre"],
+ "condition_sel": ["sel"],
+ "titrant_name": ["drug"],
+ "titrant_conc": [0.0],
+ "genotype": ["g0"],
+ "t_pre": [5.0],
+ "t_sel": [5.0],
+ "ln_cfu": [10.0],
+ "ln_cfu_std": [0.2],
+ })
+
+ captured_params = {}
+
+ def _mock_predict(orch, params, predict_sites, **kwargs):
+ captured_params["value"] = params
+ if len(predict_sites) > 1:
+ ln0 = pd.DataFrame({
+ "replicate": ["r0"],
+ "condition_pre": ["pre"],
+ "genotype": ["g0"],
+ "q0.05": [9.5],
+ "q0.5": [10.0],
+ "q0.95": [10.5],
+ })
+ return {"growth_pred": fine_df, "ln_cfu0": ln0}
+ return fine_df
+
+ from tfscreen.plot import geno_trajectory
+ with patch.object(geno_trajectory, "predict_and_plot_geno_trajectory",
+ wraps=geno_trajectory.predict_and_plot_geno_trajectory):
+ with patch("tfscreen.tfmodel.analysis.prediction.predict",
+ side_effect=_mock_predict):
+ from tfscreen.plot.geno_trajectory import predict_and_plot_geno_trajectory
+ fig = predict_and_plot_geno_trajectory(orchestrator, npz_path)
+
+ assert isinstance(captured_params["value"], dict), \
+ "npz path should be loaded to a dict before predict() is called"
+ assert "x_auto_loc" in captured_params["value"]
+ plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# presplit_df overlay
+# ---------------------------------------------------------------------------
+
+def _make_mock_orchestrator_with_presplit(presplit_df=None):
+ """Return a minimal mock orchestrator for predict_and_plot tests."""
+ from unittest.mock import MagicMock
+
+ gd = MagicMock()
+ gd.good_mask = np.ones((1, 3, 1, 1, 1, 1, 1), dtype=bool)
+ gd.t_sel = np.broadcast_to(
+ np.array([0.0, 5.0, 10.0])[None, :, None, None, None, None, None],
+ (1, 3, 1, 1, 1, 1, 1),
+ ).copy()
+ data = MagicMock()
+ data.growth = gd
+
+ orch = MagicMock()
+ orch.data = data
+ orch.growth_df = pd.DataFrame({
+ "replicate": ["r0"],
+ "condition_pre": ["pre"],
+ "condition_sel": ["sel"],
+ "titrant_name": ["drug"],
+ "titrant_conc": [0.0],
+ "genotype": ["g0"],
+ "t_pre": [5.0],
+ "t_sel": [5.0],
+ "ln_cfu": [10.0],
+ "ln_cfu_std": [0.2],
+ })
+ orch.presplit_df = presplit_df
+ return orch
+
+
+def _mock_predict_two_sites(orch, params, predict_sites, **kwargs):
+ fine_df = pd.DataFrame({
+ "replicate": ["r0"] * 5,
+ "condition_pre": ["pre"] * 5,
+ "condition_sel": ["sel"] * 5,
+ "titrant_name": ["drug"] * 5,
+ "titrant_conc": [0.0] * 5,
+ "genotype": ["g0"] * 5,
+ "t_sel": np.linspace(0, 10, 5),
+ "ln_cfu": [np.nan] * 5,
+ "ln_cfu_std": [np.nan] * 5,
+ "q0.05": np.linspace(8, 12, 5) - 0.5,
+ "q0.5": np.linspace(8, 12, 5),
+ "q0.95": np.linspace(8, 12, 5) + 0.5,
+ })
+ ln0 = pd.DataFrame({
+ "replicate": ["r0"],
+ "condition_pre": ["pre"],
+ "genotype": ["g0"],
+ "q0.05": [9.5],
+ "q0.5": [10.0],
+ "q0.95": [10.5],
+ })
+ return {"growth_pred": fine_df, "ln_cfu0": ln0}
+
+
+class TestPresplitOverlay:
+ def test_no_presplit_anchor_has_nan_ln_cfu(self):
+ """Without presplit_df the anchor row has NaN ln_cfu."""
+ from unittest.mock import patch
+ orch = _make_mock_orchestrator_with_presplit(presplit_df=None)
+
+ with patch("tfscreen.tfmodel.analysis.prediction.predict",
+ side_effect=_mock_predict_two_sites):
+ from tfscreen.plot.geno_trajectory import predict_and_plot_geno_trajectory
+ fig = predict_and_plot_geno_trajectory(orch, {"dummy": np.array(1.0)})
+
+ ax = next(a for a in fig.axes if a.get_visible())
+ # Only observed data plotted via errorbar comes from non-NaN ln_cfu.
+ # With no presplit and obs rows at t_sel>=0 only, t<0 should have no errorbar.
+ obs_lines_neg = [
+ l for l in ax.lines
+ if l.get_linewidth() != pytest.approx(1.8) # exclude predicted lines
+ and list(l.get_xdata()) != [0.0, 0.0] # exclude axvline
+ and any(x < 0 for x in l.get_xdata())
+ ]
+ assert len(obs_lines_neg) == 0
+ plt.close(fig)
+
+ def test_presplit_anchor_shows_observed_point(self):
+ """With presplit_df the anchor row has observed ln_cfu rendered as an errorbar."""
+ from unittest.mock import patch
+ presplit = pd.DataFrame({
+ "replicate": ["r0"],
+ "condition_pre": ["pre"],
+ "genotype": ["g0"],
+ "ln_cfu": [9.8],
+ "ln_cfu_std": [0.15],
+ })
+ orch = _make_mock_orchestrator_with_presplit(presplit_df=presplit)
+
+ with patch("tfscreen.tfmodel.analysis.prediction.predict",
+ side_effect=_mock_predict_two_sites):
+ from tfscreen.plot.geno_trajectory import predict_and_plot_geno_trajectory
+ fig = predict_and_plot_geno_trajectory(orch, {"dummy": np.array(1.0)})
+
+ ax = next(a for a in fig.axes if a.get_visible())
+ # errorbar lines at negative t should exist now
+ obs_lines_neg = [
+ l for l in ax.lines
+ if l.get_linewidth() != pytest.approx(1.8)
+ and list(l.get_xdata()) != [0.0, 0.0]
+ and len(l.get_xdata()) > 0
+ and any(x < 0 for x in l.get_xdata())
+ ]
+ assert len(obs_lines_neg) > 0
+ plt.close(fig)
+
+ def test_presplit_unmatched_genotype_stays_nan(self):
+ """presplit row for a different genotype does not affect anchor ln_cfu."""
+ from unittest.mock import patch
+ presplit = pd.DataFrame({
+ "replicate": ["r0"],
+ "condition_pre": ["pre"],
+ "genotype": ["other_geno"], # not "g0"
+ "ln_cfu": [9.8],
+ "ln_cfu_std": [0.15],
+ })
+ orch = _make_mock_orchestrator_with_presplit(presplit_df=presplit)
+
+ with patch("tfscreen.tfmodel.analysis.prediction.predict",
+ side_effect=_mock_predict_two_sites):
+ from tfscreen.plot.geno_trajectory import predict_and_plot_geno_trajectory
+ fig = predict_and_plot_geno_trajectory(orch, {"dummy": np.array(1.0)})
+
+ ax = next(a for a in fig.axes if a.get_visible())
+ obs_lines_neg = [
+ l for l in ax.lines
+ if l.get_linewidth() != pytest.approx(1.8)
+ and list(l.get_xdata()) != [0.0, 0.0]
+ and len(l.get_xdata()) > 0
+ and any(x < 0 for x in l.get_xdata())
+ ]
+ assert len(obs_lines_neg) == 0
+ plt.close(fig)
diff --git a/tests/tfscreen/plot/test_plot_theta_fits.py b/tests/tfscreen/plot/test_plot_theta_fits.py
index 9ec37547..89fbd214 100644
--- a/tests/tfscreen/plot/test_plot_theta_fits.py
+++ b/tests/tfscreen/plot/test_plot_theta_fits.py
@@ -27,14 +27,14 @@ def _make_df(genotypes=("wt",), titrant_names=("IPTG",),
"titrant_conc": c,
"theta_obs": rng.uniform(0.1, 0.9),
"theta_std": 0.05,
- "median": rng.uniform(0.1, 0.9),
+ "q0.5": rng.uniform(0.1, 0.9),
}
if with_std_band:
- row["lower_std"] = row["median"] - 0.1
- row["upper_std"] = row["median"] + 0.1
+ row["q0.159"] = row["q0.5"] - 0.1
+ row["q0.841"] = row["q0.5"] + 0.1
if with_95_band:
- row["lower_95"] = row["median"] - 0.2
- row["upper_95"] = row["median"] + 0.2
+ row["q0.025"] = row["q0.5"] - 0.2
+ row["q0.975"] = row["q0.5"] + 0.2
rows.append(row)
return pd.DataFrame(rows)
@@ -265,8 +265,8 @@ def test_custom_markers_accepted():
plt.close("all")
-def test_point_est_used_when_median_absent():
- """point_est column is used as centre line when median is not present."""
+def test_point_est_used_when_q0_5_absent():
+ """point_est column is used as centre line when q0.5 is not present."""
rng = np.random.default_rng(42)
df = pd.DataFrame({
"genotype": ["wt"] * 4,
@@ -282,8 +282,8 @@ def test_point_est_used_when_median_absent():
plt.close("all")
-def test_median_takes_precedence_over_point_est():
- """median is used when both median and point_est columns are present."""
+def test_q0_5_takes_precedence_over_point_est():
+ """q0.5 is used when both q0.5 and point_est columns are present."""
rng = np.random.default_rng(0)
df = pd.DataFrame({
"genotype": ["wt"] * 4,
@@ -291,7 +291,7 @@ def test_median_takes_precedence_over_point_est():
"titrant_conc": [0.1, 1.0, 10.0, 100.0],
"theta_obs": rng.uniform(0.1, 0.9, 4),
"theta_std": [0.05] * 4,
- "median": np.array([0.1, 0.2, 0.3, 0.4]),
+ "q0.5": np.array([0.1, 0.2, 0.3, 0.4]),
"point_est": np.array([0.9, 0.8, 0.7, 0.6]),
})
plt.close("all")
diff --git a/tests/tfscreen/process_raw/test_process_counts.py b/tests/tfscreen/process_raw/test_process_counts_cli.py
similarity index 100%
rename from tests/tfscreen/process_raw/test_process_counts.py
rename to tests/tfscreen/process_raw/test_process_counts_cli.py
diff --git a/tests/tfscreen/process_raw/test_process_fastq.py b/tests/tfscreen/process_raw/test_process_fastq_cli.py
similarity index 100%
rename from tests/tfscreen/process_raw/test_process_fastq.py
rename to tests/tfscreen/process_raw/test_process_fastq_cli.py
diff --git a/tests/tfscreen/simulate/test_binding_params.py b/tests/tfscreen/simulate/test_binding_params.py
new file mode 100644
index 00000000..240e1ca4
--- /dev/null
+++ b/tests/tfscreen/simulate/test_binding_params.py
@@ -0,0 +1,585 @@
+"""
+Unit tests for tfscreen.simulate.binding_params.
+"""
+
+import io
+import pytest
+import numpy as np
+import pandas as pd
+from unittest.mock import MagicMock
+
+from tfscreen.simulate.binding_params import (
+ read_binding_genotype_params,
+ build_theta_gc_override_hill_geno,
+ build_theta_gc_override_hill_mut,
+ build_binding_theta_from_params,
+ _hill_theta,
+ _fill_params_from_wt,
+ _logit,
+ _to_log_conc,
+ _wt_params_from_sim_priors,
+ SUPPORTED_COMPONENTS,
+ HILL_PARAM_COLS,
+)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def simple_csv(tmp_path):
+ """CSV with all four Hill parameter columns."""
+ content = (
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,0.01,-4.1,2.0\n"
+ "A47V,0.97,0.03,-3.8,1.8\n"
+ "K84L,0.98,0.02,-3.5,2.2\n"
+ )
+ p = tmp_path / "params.csv"
+ p.write_text(content)
+ return str(p)
+
+
+@pytest.fixture
+def partial_csv(tmp_path):
+ """CSV with only two parameter columns; others should NaN."""
+ content = (
+ "genotype,log_hill_K,hill_n\n"
+ "wt,-4.1,2.0\n"
+ "A47V,-3.8,1.8\n"
+ )
+ p = tmp_path / "partial.csv"
+ p.write_text(content)
+ return str(p)
+
+
+@pytest.fixture
+def nan_csv(tmp_path):
+ """CSV with explicit NaN values."""
+ content = (
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,0.01,-4.1,2.0\n"
+ "A47V,,,-3.8,1.8\n"
+ )
+ p = tmp_path / "nan_params.csv"
+ p.write_text(content)
+ return str(p)
+
+
+@pytest.fixture
+def wt_params():
+ return {
+ "theta_low": 0.99,
+ "theta_high": 0.01,
+ "log_hill_K": -4.1,
+ "hill_n": 2.0,
+ }
+
+
+@pytest.fixture
+def mock_sim_priors(wt_params):
+ sp = MagicMock()
+ sp.wt_theta_low = wt_params["theta_low"]
+ sp.wt_theta_high = wt_params["theta_high"]
+ sp.wt_log_K = wt_params["log_hill_K"]
+ sp.wt_hill_n = wt_params["hill_n"]
+ sp.sigma_d_logit_low = 0.3
+ sp.sigma_d_logit_delta = 0.5
+ sp.sigma_d_log_K = 0.5
+ sp.sigma_d_log_n = 0.3
+ sp.epi_tau_scale = 0.0 # no epistasis by default in tests
+ sp.epi_slab_scale = 2.0
+ sp.epi_slab_df = 4.0
+ return sp
+
+
+@pytest.fixture
+def rng():
+ return np.random.default_rng(42)
+
+
+# ---------------------------------------------------------------------------
+# read_binding_genotype_params
+# ---------------------------------------------------------------------------
+
+class TestReadBindingGenotypeParams:
+
+ def test_reads_all_params(self, simple_csv):
+ d = read_binding_genotype_params(simple_csv)
+ assert set(d.keys()) == {"wt", "A47V", "K84L"}
+ assert d["wt"]["theta_low"] == pytest.approx(0.99)
+ assert d["A47V"]["log_hill_K"] == pytest.approx(-3.8)
+
+ def test_reads_partial_columns(self, partial_csv):
+ d = read_binding_genotype_params(partial_csv)
+ assert "theta_low" not in d["wt"]
+ assert d["wt"]["log_hill_K"] == pytest.approx(-4.1)
+
+ def test_nan_preserved(self, nan_csv):
+ d = read_binding_genotype_params(nan_csv)
+ assert np.isnan(d["A47V"]["theta_low"])
+ assert np.isnan(d["A47V"]["theta_high"])
+ assert d["A47V"]["log_hill_K"] == pytest.approx(-3.8)
+
+ def test_raises_no_genotype_column(self, tmp_path):
+ p = tmp_path / "bad.csv"
+ p.write_text("theta_low,theta_high\n0.99,0.01\n")
+ with pytest.raises(ValueError, match="genotype"):
+ read_binding_genotype_params(str(p))
+
+ def test_raises_no_param_columns(self, tmp_path):
+ p = tmp_path / "bad.csv"
+ p.write_text("genotype\nwt\n")
+ with pytest.raises(ValueError, match="parameter column"):
+ read_binding_genotype_params(str(p))
+
+ def test_raises_unknown_columns(self, tmp_path):
+ p = tmp_path / "bad.csv"
+ p.write_text("genotype,theta_low,extra_col\nwt,0.99,foo\n")
+ with pytest.raises(ValueError, match="Unrecognised"):
+ read_binding_genotype_params(str(p))
+
+ def test_clips_theta_above_one(self, tmp_path):
+ """theta_low > 1 (float rounding) must be clipped and trigger a warning."""
+ p = tmp_path / "over.csv"
+ p.write_text("genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,1.000004,0.01,-4.1,2.0\n")
+ with pytest.warns(UserWarning, match="theta_low"):
+ d = read_binding_genotype_params(str(p))
+ assert d["wt"]["theta_low"] <= 1.0 - 1e-4
+
+ def test_clips_theta_below_zero(self, tmp_path):
+ """theta_high < 0 must be clipped and trigger a warning."""
+ p = tmp_path / "under.csv"
+ p.write_text("genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,-0.0001,-4.1,2.0\n")
+ with pytest.warns(UserWarning, match="theta_high"):
+ d = read_binding_genotype_params(str(p))
+ assert d["wt"]["theta_high"] >= 1e-4
+
+ def test_no_warning_for_valid_values(self, simple_csv):
+ """Valid theta values in (0, 1) must not trigger a warning."""
+ import warnings as _w
+ with _w.catch_warnings():
+ _w.simplefilter("error", UserWarning)
+ read_binding_genotype_params(simple_csv) # should not raise
+
+ def test_nan_theta_not_clipped(self, nan_csv):
+ """NaN theta values must pass through without triggering a warning."""
+ import warnings as _w
+ with _w.catch_warnings():
+ _w.simplefilter("error", UserWarning)
+ d = read_binding_genotype_params(nan_csv)
+ assert np.isnan(d["A47V"]["theta_low"])
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+class TestHillTheta:
+
+ def test_wt_curve_shape(self, wt_params):
+ log_conc = np.linspace(-10, 0, 20)
+ theta = _hill_theta(**wt_params, log_conc=log_conc)
+ assert theta.shape == (20,)
+ assert np.all(theta >= 0) and np.all(theta <= 1)
+
+ def test_decreasing_for_repressor(self, wt_params):
+ """With theta_low > theta_high, theta decreases as ligand increases."""
+ log_conc = np.linspace(-10, 0, 10)
+ theta = _hill_theta(**wt_params, log_conc=log_conc)
+ assert theta[0] > theta[-1]
+
+ def test_zero_concentration_handled(self, wt_params):
+ """Zero concentrations are passed as the sentinel; no NaN."""
+ log_conc = _to_log_conc([0.0, 0.001, 0.01])
+ theta = _hill_theta(**wt_params, log_conc=log_conc)
+ assert np.all(np.isfinite(theta))
+
+ def test_endpoints(self, wt_params):
+ """Very low conc → theta_low; very high conc → theta_high."""
+ theta_low_end = _hill_theta(**wt_params, log_conc=np.array([-30.0]))
+ theta_high_end = _hill_theta(**wt_params, log_conc=np.array([10.0]))
+ assert theta_low_end[0] == pytest.approx(wt_params["theta_low"], abs=1e-4)
+ assert theta_high_end[0] == pytest.approx(wt_params["theta_high"], abs=1e-4)
+
+
+class TestFillParamsFromWt:
+
+ def test_fills_nan_from_wt(self, wt_params):
+ params = {"theta_low": float("nan"), "log_hill_K": -3.8,
+ "theta_high": float("nan"), "hill_n": float("nan")}
+ filled = _fill_params_from_wt(params, wt_params)
+ assert filled["theta_low"] == wt_params["theta_low"]
+ assert filled["theta_high"] == wt_params["theta_high"]
+ assert filled["hill_n"] == wt_params["hill_n"]
+ assert filled["log_hill_K"] == pytest.approx(-3.8)
+
+ def test_does_not_overwrite_present(self, wt_params):
+ params = {"theta_low": 0.5, "theta_high": 0.5,
+ "log_hill_K": -3.0, "hill_n": 1.5}
+ filled = _fill_params_from_wt(params, wt_params)
+ assert filled["theta_low"] == pytest.approx(0.5)
+
+ def test_missing_key_uses_wt(self, wt_params):
+ params = {"log_hill_K": -3.8, "hill_n": 1.8}
+ filled = _fill_params_from_wt(params, wt_params)
+ assert filled["theta_low"] == wt_params["theta_low"]
+ assert filled["theta_high"] == wt_params["theta_high"]
+
+
+class TestWtParamsFromSimPriors:
+
+ def test_extracts_correct_keys(self, mock_sim_priors, wt_params):
+ result = _wt_params_from_sim_priors(mock_sim_priors)
+ assert result["theta_low"] == pytest.approx(wt_params["theta_low"])
+ assert result["log_hill_K"] == pytest.approx(wt_params["log_hill_K"])
+ assert set(result.keys()) == HILL_PARAM_COLS
+
+
+# ---------------------------------------------------------------------------
+# build_theta_gc_override_hill_geno
+# ---------------------------------------------------------------------------
+
+class TestBuildThetaGcOverrideHillGeno:
+
+ def test_returns_correct_genotypes(self, simple_csv, wt_params):
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.linspace(-10, 0, 5)
+ override, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ assert set(override.keys()) == {"wt", "A47V", "K84L"}
+
+ def test_theta_shape(self, simple_csv, wt_params):
+ params_dict = read_binding_genotype_params(simple_csv)
+ C = 8
+ log_conc = np.linspace(-10, 0, C)
+ override, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ for g, theta in override.items():
+ assert theta.shape == (C,), f"Shape mismatch for {g}"
+
+ def test_wt_uses_wt_params(self, simple_csv, wt_params):
+ """WT row in CSV should produce the exact WT curve."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.linspace(-10, 0, 20)
+ override, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ expected = _hill_theta(**wt_params, log_conc=log_conc)
+ np.testing.assert_allclose(override["wt"], expected, rtol=1e-6)
+
+ def test_mutant_uses_measured_params(self, simple_csv, wt_params):
+ """Non-WT row should use the measured parameters, not WT."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.linspace(-10, 0, 20)
+ override, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ wt_curve = _hill_theta(**wt_params, log_conc=log_conc)
+ # A47V has a different log_hill_K, so its curve should differ from WT
+ assert not np.allclose(override["A47V"], wt_curve)
+
+ def test_nan_params_filled_from_wt(self, nan_csv, wt_params):
+ """NaN theta_low / theta_high should fall back to WT values."""
+ params_dict = read_binding_genotype_params(nan_csv)
+ log_conc = np.linspace(-10, 0, 5)
+ override, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ # A47V has NaN theta_low / theta_high → WT values used
+ # its log_hill_K is -3.8 (different from WT -4.1), so not identical to WT
+ assert override["A47V"].shape == (5,)
+ assert np.all(np.isfinite(override["A47V"]))
+
+ def test_theta_in_unit_interval(self, simple_csv, wt_params):
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.linspace(-15, 5, 30)
+ override, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ for theta in override.values():
+ assert np.all(theta >= 0) and np.all(theta <= 1)
+
+ def test_effective_params_returned(self, simple_csv, wt_params):
+ """effective_params matches the CSV values (NaN filled from wt_params)."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.linspace(-10, 0, 5)
+ _, eff = build_theta_gc_override_hill_geno(params_dict, log_conc, wt_params)
+ assert set(eff.keys()) == {"wt", "A47V", "K84L"}
+ assert eff["wt"]["theta_low"] == pytest.approx(0.99)
+ assert eff["A47V"]["log_hill_K"] == pytest.approx(-3.8)
+ assert eff["K84L"]["hill_n"] == pytest.approx(2.2)
+
+
+# ---------------------------------------------------------------------------
+# build_theta_gc_override_hill_mut
+# ---------------------------------------------------------------------------
+
+def _make_sim_data_for_genotypes(genotypes):
+ """Build a real SimData for a small genotype list (no thermo data)."""
+ import pandas as pd
+ from tfscreen.simulate.sim_data_class import build_sim_data
+
+ library_df = pd.DataFrame({"genotype": genotypes})
+ sample_df = pd.DataFrame({"titrant_conc": [0.0, 0.001, 0.01, 0.1, 1.0]})
+ return build_sim_data(library_df, sample_df)
+
+
+class TestBuildThetaGcOverrideHillMut:
+
+ @pytest.fixture
+ def small_genotypes(self):
+ return ["wt", "A47V", "K84L", "A47V/K84L"]
+
+ @pytest.fixture
+ def small_sim_data(self, small_genotypes):
+ return _make_sim_data_for_genotypes(small_genotypes)
+
+ def test_returns_all_genotypes(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ assert set(result.keys()) == set(small_genotypes)
+
+ def test_theta_shape(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ C = len(log_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ for g, theta in result.items():
+ assert theta.shape == (C,), f"Shape mismatch for {g}"
+
+ def test_theta_in_unit_interval(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ for g, theta in result.items():
+ assert np.all(theta >= 0) and np.all(theta <= 1), f"Out of [0,1] for {g}"
+
+ def test_wt_receives_exact_wt_curve(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ """WT genotype has no mutations so its assembled curve = WT reference curve."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ # WT from CSV: theta_low=0.99, theta_high=0.01, log_hill_K=-4.1, hill_n=2.0
+ expected_wt = _hill_theta(0.99, 0.01, -4.1, 2.0, log_conc)
+ np.testing.assert_allclose(result["wt"], expected_wt, rtol=1e-5)
+
+ def test_single_mutant_uses_measured_params(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ """A single-mutant entry in the CSV should produce its measured curve exactly."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ expected_a47v = _hill_theta(0.97, 0.03, -3.8, 1.8, log_conc)
+ np.testing.assert_allclose(result["A47V"], expected_a47v, rtol=1e-5)
+
+ def test_double_differs_from_wt(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ """Double mutant should differ from WT (deltas from both singles applied)."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ expected_wt = _hill_theta(0.99, 0.01, -4.1, 2.0, log_conc)
+ assert not np.allclose(result["A47V/K84L"], expected_wt)
+
+ def test_unmeasured_mut_uses_prior(self, wt_params, mock_sim_priors, rng):
+ """A mutation not in the CSV gets a random delta from the prior."""
+ genotypes = ["wt", "M99A"] # M99A not in CSV
+ sim_data = _make_sim_data_for_genotypes(genotypes)
+ params_dict = {"wt": wt_params} # only WT in CSV
+ log_conc = np.array(sim_data.log_titrant_conc)
+
+ r1, _ = build_theta_gc_override_hill_mut(
+ params_dict, genotypes, sim_data, mock_sim_priors,
+ log_conc, np.random.default_rng(1),
+ )
+ r2, _ = build_theta_gc_override_hill_mut(
+ params_dict, genotypes, sim_data, mock_sim_priors,
+ log_conc, np.random.default_rng(2),
+ )
+ # WT is pinned: same both times
+ np.testing.assert_allclose(r1["wt"], r2["wt"])
+ # M99A is random: different seeds → different curves
+ assert not np.allclose(r1["M99A"], r2["M99A"])
+
+ def test_csv_wt_overrides_sim_priors(self, small_genotypes, small_sim_data,
+ mock_sim_priors, rng, tmp_path):
+ """When CSV has a 'wt' row, that WT reference overrides sim_priors."""
+ # Use a WT with very different log_hill_K
+ content = (
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.95,0.05,-6.0,1.5\n" # dramatically different from sim_priors WT
+ )
+ p = tmp_path / "alt_wt.csv"
+ p.write_text(content)
+ params_dict = read_binding_genotype_params(str(p))
+
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ expected_wt = _hill_theta(0.95, 0.05, -6.0, 1.5, log_conc)
+ np.testing.assert_allclose(result["wt"], expected_wt, rtol=1e-5)
+
+ def test_direct_double_in_csv_overrides_assembled(self, small_genotypes,
+ small_sim_data, mock_sim_priors,
+ rng, tmp_path):
+ """A directly-measured double mutant in the CSV overrides the assembled value."""
+ content = (
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,0.01,-4.1,2.0\n"
+ "A47V/K84L,0.90,0.10,-5.0,3.0\n" # direct measurement for the double
+ )
+ p = tmp_path / "with_double.csv"
+ p.write_text(content)
+ params_dict = read_binding_genotype_params(str(p))
+
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ result, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ expected_double = _hill_theta(0.90, 0.10, -5.0, 3.0, log_conc)
+ np.testing.assert_allclose(result["A47V/K84L"], expected_double, rtol=1e-5)
+
+ def test_effective_params_all_genotypes(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ """effective_params covers all library genotypes."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ _, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ assert set(eff.keys()) == set(small_genotypes)
+ for g, p in eff.items():
+ assert set(p.keys()) == {"theta_low", "theta_high", "log_hill_K", "hill_n"}
+
+ def test_effective_params_measured_singles(self, simple_csv, small_genotypes,
+ small_sim_data, mock_sim_priors, rng):
+ """effective_params for CSV-measured singles matches their CSV values."""
+ params_dict = read_binding_genotype_params(simple_csv)
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ _, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ assert eff["wt"]["theta_low"] == pytest.approx(0.99, abs=1e-4)
+ assert eff["wt"]["theta_high"] == pytest.approx(0.01, abs=1e-4)
+ assert eff["A47V"]["log_hill_K"] == pytest.approx(-3.8, abs=1e-4)
+ assert eff["K84L"]["hill_n"] == pytest.approx(2.2, abs=1e-4)
+
+ def test_effective_params_direct_double(self, small_genotypes, small_sim_data,
+ mock_sim_priors, rng, tmp_path):
+ """effective_params for a directly-measured double uses the CSV values."""
+ content = (
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,0.01,-4.1,2.0\n"
+ "A47V/K84L,0.90,0.10,-5.0,3.0\n"
+ )
+ p = tmp_path / "with_double.csv"
+ p.write_text(content)
+ params_dict = read_binding_genotype_params(str(p))
+ log_conc = np.array(small_sim_data.log_titrant_conc)
+ _, eff = build_theta_gc_override_hill_mut(
+ params_dict, small_genotypes, small_sim_data,
+ mock_sim_priors, log_conc, rng,
+ )
+ assert eff["A47V/K84L"]["theta_low"] == pytest.approx(0.90, abs=1e-4)
+ assert eff["A47V/K84L"]["log_hill_K"] == pytest.approx(-5.0, abs=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# build_binding_theta_from_params
+# ---------------------------------------------------------------------------
+
+class TestBuildBindingThetaFromParams:
+
+ @pytest.fixture
+ def params_dict(self, simple_csv):
+ return read_binding_genotype_params(simple_csv)
+
+ def test_output_shape(self, params_dict, wt_params, rng):
+ binding_concs = [0.0, 0.001, 0.01, 0.1, 1.0]
+ df = build_binding_theta_from_params(
+ params_dict, binding_concs, "iptg", noise=0.0, rng=rng, wt_params=wt_params
+ )
+ assert set(df.columns) == {"genotype", "titrant_name", "titrant_conc", "theta_true"}
+ # 3 genotypes × 5 concentrations = 15 rows
+ assert len(df) == 3 * len(binding_concs)
+
+ def test_correct_titrant_name(self, params_dict, wt_params, rng):
+ df = build_binding_theta_from_params(
+ params_dict, [0.01], "iptg", noise=0.0, rng=rng, wt_params=wt_params
+ )
+ assert (df["titrant_name"] == "iptg").all()
+
+ def test_no_noise(self, params_dict, wt_params, rng):
+ """With noise=0, theta_true should exactly match Hill equation output."""
+ binding_concs = [0.0, 0.001, 0.01]
+ log_concs = _to_log_conc(binding_concs)
+ df = build_binding_theta_from_params(
+ params_dict, binding_concs, "iptg", noise=0.0, rng=rng, wt_params=wt_params
+ )
+ wt_rows = df[df["genotype"] == "wt"].sort_values("titrant_conc")
+ expected = _hill_theta(0.99, 0.01, -4.1, 2.0, log_concs)
+ np.testing.assert_allclose(wt_rows["theta_true"].values, expected, rtol=1e-6)
+
+ def test_noise_changes_values(self, params_dict, wt_params):
+ """With noise > 0, two runs with different seeds produce different theta_true."""
+ binding_concs = [0.0, 0.001, 0.01, 0.1]
+ df1 = build_binding_theta_from_params(
+ params_dict, binding_concs, "iptg", noise=0.05,
+ rng=np.random.default_rng(1), wt_params=wt_params,
+ )
+ df2 = build_binding_theta_from_params(
+ params_dict, binding_concs, "iptg", noise=0.05,
+ rng=np.random.default_rng(2), wt_params=wt_params,
+ )
+ assert not np.allclose(df1["theta_true"].values, df2["theta_true"].values)
+
+ def test_theta_in_unit_interval(self, params_dict, wt_params):
+ """theta_true must remain in [0, 1] even with noise."""
+ binding_concs = np.linspace(0.0, 1.0, 20)
+ df = build_binding_theta_from_params(
+ params_dict, binding_concs, "iptg", noise=0.1,
+ rng=np.random.default_rng(99), wt_params=wt_params,
+ )
+ assert (df["theta_true"] >= 0).all() and (df["theta_true"] <= 1).all()
+
+ def test_nan_filled_from_wt(self, nan_csv, wt_params, rng):
+ """Genotypes with NaN params should produce finite theta using WT fallback."""
+ params_dict = read_binding_genotype_params(nan_csv)
+ df = build_binding_theta_from_params(
+ params_dict, [0.0, 0.01, 1.0], "iptg", noise=0.0, rng=rng, wt_params=wt_params,
+ )
+ assert np.all(np.isfinite(df["theta_true"].values))
+
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+def test_supported_components():
+ assert "hill_geno" in SUPPORTED_COMPONENTS
+ assert "hill_mut" in SUPPORTED_COMPONENTS
+
+def test_hill_param_cols():
+ assert HILL_PARAM_COLS == {"theta_low", "theta_high", "log_hill_K", "hill_n"}
diff --git a/tests/tfscreen/simulate/test_generate_presplit_data.py b/tests/tfscreen/simulate/test_generate_presplit_data.py
new file mode 100644
index 00000000..83d3cf7c
--- /dev/null
+++ b/tests/tfscreen/simulate/test_generate_presplit_data.py
@@ -0,0 +1,245 @@
+"""
+Tests for _generate_presplit_data in simulate_cli.
+
+Coverage:
+ - Output columns and shape for a simple two-genotype, two-condition case
+ - Genotypes absent from the transformation pool (ln_cfu_0 = -inf) → NaN ln_cfu
+ - extra_noise parameter inflates ln_cfu_std
+ - Output is deterministic given the same rng seed
+ - run_simulation_from_config writes a presplit CSV when presplit_data is in cf
+"""
+import pytest
+import numpy as np
+import pandas as pd
+from unittest.mock import patch, MagicMock
+
+from tfscreen.simulate.scripts.simulate_cli import _generate_presplit_data
+
+
+# ---------------------------------------------------------------------------
+# Minimal fixtures
+# ---------------------------------------------------------------------------
+
+def _make_inputs(
+ replicates=(1, 2),
+ condition_pres=("kanR", "pheS"),
+ genotypes=("wt", "A1V", "A2V"),
+ ln_cfu_0_val=10.0,
+):
+ """Build minimal combined_sample_df and combined_counts_df."""
+ sample_rows = []
+ counts_rows = []
+ sample_id = 0
+
+ for rep in replicates:
+ for cp in condition_pres:
+ # Two selection samples per (rep, cp) (different t_sel)
+ for t_sel in (60.0, 90.0):
+ sample_rows.append(
+ {"sample": sample_id, "replicate": rep,
+ "condition_pre": cp, "t_sel": t_sel,
+ "sample_cfu": 1e8, "sample_cfu_std": 5e6}
+ )
+ for geno in genotypes:
+ counts_rows.append(
+ {"sample": sample_id, "genotype": geno,
+ "ln_cfu_0": ln_cfu_0_val, "counts": 100}
+ )
+ sample_id += 1
+
+ sample_df = pd.DataFrame(sample_rows).set_index("sample")
+ counts_df = pd.DataFrame(counts_rows)
+ return sample_df, counts_df
+
+
+def _minimal_cf(noise=0.0, extra_keys=None):
+ cf = {
+ "cfu0": 1e8,
+ "total_num_reads": 30_000_000,
+ "prob_index_hop": None,
+ "presplit_data": {"noise": noise},
+ }
+ if extra_keys:
+ cf.update(extra_keys)
+ return cf
+
+
+# ---------------------------------------------------------------------------
+# Basic output structure
+# ---------------------------------------------------------------------------
+
+def test_output_columns():
+ sample_df, counts_df = _make_inputs()
+ result = _generate_presplit_data(sample_df, counts_df, _minimal_cf(),
+ np.random.default_rng(0))
+ for col in ["replicate", "condition_pre", "genotype", "ln_cfu", "ln_cfu_std",
+ "ln_cfu_0_true"]:
+ assert col in result.columns, f"missing column: {col}"
+
+
+def test_output_row_count():
+ """One row per (replicate, condition_pre, genotype)."""
+ replicates = (1, 2)
+ cps = ("kanR", "pheS")
+ genos = ("wt", "A1V", "A2V")
+ sample_df, counts_df = _make_inputs(replicates=replicates,
+ condition_pres=cps,
+ genotypes=genos)
+ result = _generate_presplit_data(sample_df, counts_df, _minimal_cf(),
+ np.random.default_rng(0))
+ expected = len(replicates) * len(cps) * len(genos)
+ assert len(result) == expected
+
+
+def test_ln_cfu_is_finite_for_present_genotypes():
+ sample_df, counts_df = _make_inputs()
+ result = _generate_presplit_data(sample_df, counts_df, _minimal_cf(),
+ np.random.default_rng(0))
+ # All genotypes have finite ln_cfu_0 so ln_cfu should be finite
+ assert result["ln_cfu"].notna().all()
+
+
+def test_ln_cfu_std_is_positive():
+ sample_df, counts_df = _make_inputs()
+ result = _generate_presplit_data(sample_df, counts_df, _minimal_cf(),
+ np.random.default_rng(0))
+ assert (result["ln_cfu_std"] > 0).all()
+
+
+# ---------------------------------------------------------------------------
+# Absent genotypes (ln_cfu_0 = -inf)
+# ---------------------------------------------------------------------------
+
+def test_absent_genotype_lower_lncfu():
+ """A genotype with ln_cfu_0 = -inf (never transformed) gets zero counts
+ from the multinomial draw. After applying the pseudocount it still gets a
+ finite (but very low) ln_cfu estimate — the same behaviour as
+ counts_to_lncfu — and its ln_cfu is substantially lower than a present
+ genotype."""
+ sample_df, counts_df = _make_inputs()
+ # Set one genotype's ln_cfu_0 to -inf across all samples
+ counts_df.loc[counts_df["genotype"] == "A2V", "ln_cfu_0"] = -np.inf
+ result = _generate_presplit_data(sample_df, counts_df, _minimal_cf(),
+ np.random.default_rng(0))
+ absent_ln = result.loc[result["genotype"] == "A2V", "ln_cfu"].mean()
+ present_ln = result.loc[result["genotype"] == "wt", "ln_cfu"].mean()
+ # Absent genotype should be much lower than a normally-present genotype
+ assert absent_ln < present_ln - 5.0
+
+
+# ---------------------------------------------------------------------------
+# extra_noise parameter
+# ---------------------------------------------------------------------------
+
+def test_extra_noise_inflates_ln_cfu_std():
+ """Adding extra noise should increase the mean ln_cfu_std."""
+ sample_df, counts_df = _make_inputs()
+ rng0 = np.random.default_rng(42)
+ rng1 = np.random.default_rng(42)
+
+ result_no_noise = _generate_presplit_data(sample_df, counts_df,
+ _minimal_cf(noise=0.0), rng0)
+ result_with_noise = _generate_presplit_data(sample_df, counts_df,
+ _minimal_cf(noise=0.5), rng1)
+
+ mean_std_no_noise = result_no_noise["ln_cfu_std"].mean()
+ mean_std_with_noise = result_with_noise["ln_cfu_std"].mean()
+ assert mean_std_with_noise > mean_std_no_noise
+
+
+# ---------------------------------------------------------------------------
+# Reproducibility
+# ---------------------------------------------------------------------------
+
+def test_reproducibility():
+ """Same rng seed → identical output."""
+ sample_df, counts_df = _make_inputs()
+ r1 = _generate_presplit_data(sample_df.copy(), counts_df.copy(), _minimal_cf(),
+ np.random.default_rng(7))
+ r2 = _generate_presplit_data(sample_df.copy(), counts_df.copy(), _minimal_cf(),
+ np.random.default_rng(7))
+ pd.testing.assert_frame_equal(r1.reset_index(drop=True),
+ r2.reset_index(drop=True))
+
+
+# ---------------------------------------------------------------------------
+# Integration: run_simulation_from_config writes presplit CSV
+# ---------------------------------------------------------------------------
+
+def test_run_simulation_writes_presplit_csv(tmp_path):
+ """presplit CSV is written when presplit_data block is in the config."""
+ from tfscreen.simulate.scripts.simulate_cli import run_simulation_from_config
+
+ lib_df = pd.DataFrame({"genotype": ["wt", "A1V"]})
+ pheno_df = pd.DataFrame({"genotype": ["wt", "A1V"]})
+ theta_df = pd.DataFrame({"genotype": ["wt", "A1V"]})
+ params_df = pd.DataFrame({"genotype": ["wt", "A1V"],
+ "dk_geno": [0.0, 0.0], "activity": [1.0, 1.0]})
+
+ # Sample/counts for one (replicate, condition_pre, t_sel) combo
+ # Real _simulate_library_group returns sample_df with "sample" as a
+ # regular column and an unnamed integer index.
+ sample_df = pd.DataFrame([{
+ "sample": 0, "replicate": 1, "condition_pre": "kanR",
+ "t_sel": 60.0, "sample_cfu": 1e8, "sample_cfu_std": 5e6,
+ }], index=[0])
+ counts_df = pd.DataFrame([
+ {"sample": 0, "genotype": "wt", "counts": 500, "ln_cfu_0": 10.0},
+ {"sample": 0, "genotype": "A1V", "counts": 500, "ln_cfu_0": 10.0},
+ ])
+ growth_df = pd.DataFrame({"genotype": ["wt", "A1V"], "ln_cfu": [10.0, 9.9]})
+
+ cf = {
+ "seed": 1,
+ "cfu0": 1e8,
+ "total_num_reads": 10_000_000,
+ "prob_index_hop": None,
+ "presplit_data": {"noise": 0.0},
+ }
+
+ with patch("tfscreen.util.read_yaml", return_value=cf), \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.selection_experiment",
+ return_value=(sample_df, counts_df)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.counts_to_lncfu",
+ return_value=growth_df):
+ run_simulation_from_config("fake_config.yaml", str(tmp_path))
+
+ presplit_path = tmp_path / "tfs_sim_presplit.csv"
+ assert presplit_path.exists(), "presplit CSV was not written"
+ presplit_df = pd.read_csv(presplit_path)
+ for col in ["replicate", "condition_pre", "genotype", "ln_cfu", "ln_cfu_std"]:
+ assert col in presplit_df.columns
+
+
+def test_run_simulation_no_presplit_without_config(tmp_path):
+ """presplit CSV is NOT written when presplit_data is absent from config."""
+ from tfscreen.simulate.scripts.simulate_cli import run_simulation_from_config
+
+ lib_df = pd.DataFrame({"genotype": ["wt"]})
+ pheno_df = pd.DataFrame({"genotype": ["wt"]})
+ theta_df = pd.DataFrame({"genotype": ["wt"]})
+ params_df = pd.DataFrame({"genotype": ["wt"], "dk_geno": [0.0], "activity": [1.0]})
+ sample_df = pd.DataFrame([{"sample": 0, "replicate": 1,
+ "condition_pre": "kanR", "t_sel": 60.0,
+ "sample_cfu": 1e8, "sample_cfu_std": 5e6}],
+ index=[0])
+ counts_df = pd.DataFrame([{"sample": 0, "genotype": "wt",
+ "counts": 1000, "ln_cfu_0": 10.0}])
+ growth_df = pd.DataFrame({"genotype": ["wt"], "ln_cfu": [10.0]})
+
+ cf = {"seed": 1, "cfu0": 1e8, "total_num_reads": 1_000_000,
+ "prob_index_hop": None}
+
+ with patch("tfscreen.util.read_yaml", return_value=cf), \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.selection_experiment",
+ return_value=(sample_df, counts_df)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.counts_to_lncfu",
+ return_value=growth_df):
+ run_simulation_from_config("fake_config.yaml", str(tmp_path))
+
+ presplit_path = tmp_path / "tfs_sim_presplit.csv"
+ assert not presplit_path.exists(), "presplit CSV should not be written without config block"
diff --git a/tests/tfscreen/simulate/test_library_prediction.py b/tests/tfscreen/simulate/test_library_prediction.py
index 7407b5b7..464c8bd0 100644
--- a/tests/tfscreen/simulate/test_library_prediction.py
+++ b/tests/tfscreen/simulate/test_library_prediction.py
@@ -10,7 +10,7 @@ def mock_config():
return {
"condition_blocks": [{"some": "block"}],
"theta_component": "mock_theta",
- "theta_rng_seed": 7,
+ "seed": 7,
"thermo_data": None,
"theta_priors": None,
"growth": {"cond_A": {"m": 1.0, "b": 0.0}},
@@ -43,15 +43,18 @@ def test_library_prediction_success(mocker, mock_config):
mock_phenotype_df = pd.DataFrame({"theta": [0.5, 0.3]})
mock_genotype_theta_df = pd.DataFrame({"genotype": ["wt", "M1"]})
+ mock_parameters_df = pd.DataFrame({"genotype": ["wt", "M1"],
+ "dk_geno": [0.0, -0.01],
+ "activity": [1.0, 1.0]})
mock_thermo = mocker.patch(
"tfscreen.simulate.library_prediction.thermo_to_growth",
- return_value=(mock_phenotype_df, mock_genotype_theta_df),
+ return_value=(mock_phenotype_df, mock_genotype_theta_df, mock_parameters_df),
)
mock_jax_key = mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey",
return_value="mock_key")
- lib_df, pheno_df, theta_df = library_prediction(cf="config.yaml")
+ lib_df, pheno_df, theta_df, params_df, binding_df = library_prediction(cf="config.yaml")
tfscreen.util.read_yaml.assert_called_once_with("config.yaml", override_keys=None)
mock_lm_cls.assert_called_once_with(mock_config)
@@ -61,7 +64,7 @@ def test_library_prediction_success(mocker, mock_config):
sample_df=mock_sample_df,
thermo_data=None,
)
- mock_jax_key.assert_called_once_with(7)
+ mock_jax_key.assert_called_once_with(7) # seed value
mock_thermo.assert_called_once()
_, kwargs = mock_thermo.call_args
@@ -73,12 +76,417 @@ def test_library_prediction_success(mocker, mock_config):
assert kwargs["dk_geno_hyper_shift"] == mock_config["dk_geno_hyper_shift"]
assert kwargs["activity_component"] == "fixed" # default when key absent
assert kwargs["activity_priors_overrides"] is None
+ assert kwargs["rng"] is not None # seeded RNG passed through
assert lib_df.equals(mock_library_df)
assert pheno_df.equals(mock_phenotype_df)
assert theta_df.equals(mock_genotype_theta_df)
+ assert params_df.equals(mock_parameters_df)
+ assert binding_df is None
def test_library_prediction_config_error():
with pytest.raises(FileNotFoundError):
library_prediction("nonexistent_config.yaml")
+
+
+def test_library_prediction_dk_geno_zero_passes_flag(mocker, mock_config):
+ """dk_geno_zero=True in config must be forwarded to thermo_to_growth."""
+ config = {**mock_config, "dk_geno_zero": True}
+ mocker.patch("tfscreen.util.read_yaml", return_value=config)
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.library_manager.LibraryManager"
+ ).return_value.build_library_df.return_value = pd.DataFrame({"genotype": ["wt", "M1"]})
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sample_dataframes",
+ return_value=pd.DataFrame({"titrant_conc": [0.0]}),
+ )
+ mocker.patch("tfscreen.simulate.library_prediction.build_sim_data", return_value=MagicMock())
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+
+ mock_thermo = mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ return_value=(
+ pd.DataFrame({"theta": [0.5]}),
+ pd.DataFrame({"genotype": ["wt"]}),
+ pd.DataFrame({"genotype": ["wt"], "dk_geno": [0.0], "activity": [1.0]}),
+ ),
+ )
+
+ library_prediction(cf="config.yaml")
+
+ _, kwargs = mock_thermo.call_args
+ assert kwargs["dk_geno_zero"] is True
+
+
+def test_library_prediction_dk_geno_zero_makes_hyper_params_optional(mocker):
+ """When dk_geno_zero=True the three dk_geno_hyper_* keys may be absent."""
+ config = {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "mock_theta",
+ "seed": 0,
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ "dk_geno_zero": True,
+ # dk_geno_hyper_* intentionally absent
+ }
+ mocker.patch("tfscreen.util.read_yaml", return_value=config)
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.library_manager.LibraryManager"
+ ).return_value.build_library_df.return_value = pd.DataFrame({"genotype": ["wt"]})
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sample_dataframes",
+ return_value=pd.DataFrame({"titrant_conc": [0.0]}),
+ )
+ mocker.patch("tfscreen.simulate.library_prediction.build_sim_data", return_value=MagicMock())
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ return_value=(
+ pd.DataFrame({"theta": [0.5]}),
+ pd.DataFrame({"genotype": ["wt"]}),
+ pd.DataFrame({"genotype": ["wt"], "dk_geno": [0.0], "activity": [1.0]}),
+ ),
+ )
+
+ # Must not raise KeyError despite missing dk_geno_hyper_* keys
+ library_prediction(cf="config.yaml")
+
+
+def test_library_prediction_missing_hyper_params_raises_without_dk_geno_zero(mocker):
+ """Without dk_geno_zero, missing dk_geno_hyper_* must raise KeyError."""
+ config = {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "mock_theta",
+ "seed": 0,
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ # dk_geno_hyper_* absent and dk_geno_zero not set
+ }
+ mocker.patch("tfscreen.util.read_yaml", return_value=config)
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.library_manager.LibraryManager"
+ ).return_value.build_library_df.return_value = pd.DataFrame({"genotype": ["wt"]})
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sample_dataframes",
+ return_value=pd.DataFrame({"titrant_conc": [0.0]}),
+ )
+ mocker.patch("tfscreen.simulate.library_prediction.build_sim_data", return_value=MagicMock())
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+
+ with pytest.raises(KeyError):
+ library_prediction(cf="config.yaml")
+
+
+# ---------------------------------------------------------------------------
+# Unified seed: seed drives both JAX key and numpy RNG
+# ---------------------------------------------------------------------------
+
+def _make_mock_thermo_return():
+ return (
+ pd.DataFrame({"theta": [0.5]}),
+ pd.DataFrame({"genotype": ["wt"]}),
+ pd.DataFrame({"genotype": ["wt"], "dk_geno": [0.0], "activity": [1.0]}),
+ )
+
+
+def _patch_library_deps(mocker, config):
+ mocker.patch("tfscreen.util.read_yaml", return_value=config)
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.library_manager.LibraryManager"
+ ).return_value.build_library_df.return_value = pd.DataFrame({"genotype": ["wt"]})
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sample_dataframes",
+ return_value=pd.DataFrame({"titrant_conc": [0.0]}),
+ )
+ mocker.patch("tfscreen.simulate.library_prediction.build_sim_data", return_value=MagicMock())
+
+
+def test_jax_key_derived_from_seed(mocker):
+ """JAX PRNGKey is called with the value of seed."""
+ config = {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "mock_theta",
+ "seed": 42,
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ "dk_geno_hyper_loc": -3.5,
+ "dk_geno_hyper_scale": 1.0,
+ "dk_geno_hyper_shift": 0.02,
+ }
+ _patch_library_deps(mocker, config)
+ mock_jax_key = mocker.patch(
+ "tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k"
+ )
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ return_value=_make_mock_thermo_return(),
+ )
+
+ library_prediction(cf="config.yaml")
+
+ mock_jax_key.assert_called_once_with(42)
+
+
+def test_jax_key_defaults_to_zero_when_no_seed(mocker):
+ """When seed is absent the JAX PRNGKey falls back to 0."""
+ config = {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "mock_theta",
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ "dk_geno_hyper_loc": -3.5,
+ "dk_geno_hyper_scale": 1.0,
+ "dk_geno_hyper_shift": 0.02,
+ }
+ _patch_library_deps(mocker, config)
+ mock_jax_key = mocker.patch(
+ "tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k"
+ )
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ return_value=_make_mock_thermo_return(),
+ )
+
+ library_prediction(cf="config.yaml")
+
+ mock_jax_key.assert_called_once_with(0)
+
+
+def test_numpy_rng_passed_to_thermo_to_growth(mocker):
+ """A seeded numpy RNG is always forwarded to thermo_to_growth."""
+ import numpy as np
+ config = {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "mock_theta",
+ "seed": 99,
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ "dk_geno_hyper_loc": -3.5,
+ "dk_geno_hyper_scale": 1.0,
+ "dk_geno_hyper_shift": 0.02,
+ }
+ _patch_library_deps(mocker, config)
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+ mock_thermo = mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ return_value=_make_mock_thermo_return(),
+ )
+
+ library_prediction(cf="config.yaml")
+
+ _, kwargs = mock_thermo.call_args
+ assert isinstance(kwargs["rng"], np.random.Generator)
+
+
+def test_numpy_rng_seeded_with_seed(mocker):
+ """The numpy RNG passed to thermo_to_growth is seeded from seed,
+ so two calls with the same seed produce the same first draw."""
+ import numpy as np
+ captured_rngs = []
+
+ def capture_rng(*args, **kwargs):
+ captured_rngs.append(kwargs["rng"])
+ return _make_mock_thermo_return()
+
+ config = {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "mock_theta",
+ "seed": 5,
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ "dk_geno_hyper_loc": -3.5,
+ "dk_geno_hyper_scale": 1.0,
+ "dk_geno_hyper_shift": 0.02,
+ }
+
+ for _ in range(2):
+ _patch_library_deps(mocker, config)
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ side_effect=capture_rng,
+ )
+ library_prediction(cf="config.yaml")
+
+ draw_a = captured_rngs[0].integers(0, 2**32)
+ draw_b = captured_rngs[1].integers(0, 2**32)
+ assert draw_a == draw_b, "Same seed must produce identical RNG state"
+
+
+# ---------------------------------------------------------------------------
+# genotype_params_file integration
+# ---------------------------------------------------------------------------
+
+def _make_params_csv(tmp_path, content):
+ p = tmp_path / "params.csv"
+ p.write_text(content)
+ return str(p)
+
+
+def _base_config_with_binding(params_path, genotypes=None):
+ binding = {
+ "titrant_name": "iptg",
+ "titrant_conc": [0.0, 0.001, 0.01],
+ "noise": 0.0,
+ "genotype_params_file": params_path,
+ }
+ if genotypes is not None:
+ binding["genotypes"] = genotypes
+ return {
+ "condition_blocks": [{"some": "block"}],
+ "theta_component": "hill_geno",
+ "seed": 7,
+ "thermo_data": None,
+ "theta_priors": None,
+ "growth": {"cond_A": {"m": 1.0, "b": 0.0}},
+ "dk_geno_hyper_loc": -3.5,
+ "dk_geno_hyper_scale": 1.0,
+ "dk_geno_hyper_shift": 0.02,
+ "binding_data": binding,
+ }
+
+
+def _patch_for_params_file(mocker, config):
+ """Patch all heavy deps; returns (mock_thermo, mock_sim_data)."""
+ mocker.patch("tfscreen.util.read_yaml", return_value=config)
+
+ mock_library_df = pd.DataFrame({"genotype": ["wt", "A47V"]})
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.library_manager.LibraryManager"
+ ).return_value.build_library_df.return_value = mock_library_df
+
+ import numpy as np
+ import jax.numpy as jnp
+ mock_sim_data = MagicMock()
+ mock_sim_data.log_titrant_conc = jnp.array(
+ [np.log(1e-20), np.log(0.001), np.log(0.01)]
+ )
+ mock_sim_data.num_mutation = 1
+ mock_sim_data.num_pair = 0
+ mock_sim_data.mut_nnz_mut_idx = np.array([0], dtype=np.int32)
+ mock_sim_data.mut_nnz_geno_idx = np.array([1], dtype=np.int32)
+ mock_sim_data.pair_nnz_pair_idx = None
+ mock_sim_data.pair_nnz_geno_idx = None
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sample_dataframes",
+ return_value=pd.DataFrame({"titrant_conc": [0.0, 0.001, 0.01]}),
+ )
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sim_data",
+ return_value=mock_sim_data,
+ )
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+
+ mock_thermo = mocker.patch(
+ "tfscreen.simulate.library_prediction.thermo_to_growth",
+ return_value=(
+ pd.DataFrame({"theta": [0.5, 0.3]}),
+ pd.DataFrame({"genotype": ["wt", "A47V"]}),
+ pd.DataFrame({
+ "genotype": ["wt", "A47V"],
+ "dk_geno": [0.0, -0.01],
+ "activity": [1.0, 1.0],
+ }),
+ ),
+ )
+ return mock_thermo, mock_sim_data
+
+
+def test_params_file_produces_binding_theta_df(mocker, tmp_path):
+ """When genotype_params_file is set, binding_theta_df is populated."""
+ csv = _make_params_csv(
+ tmp_path,
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,0.01,-4.1,2.0\n"
+ "A47V,0.97,0.03,-3.8,1.8\n",
+ )
+ config = _base_config_with_binding(csv)
+ _patch_for_params_file(mocker, config)
+
+ # hill_geno module must be importable for sim_priors
+ from tfscreen.tfmodel.generative.registry import model_registry
+ assert "hill_geno" in model_registry["theta"]
+
+ _, _, _, _, binding_df = library_prediction(cf="config.yaml")
+
+ assert binding_df is not None
+ assert set(binding_df["genotype"].unique()) == {"wt", "A47V"}
+ assert set(binding_df.columns) >= {"genotype", "titrant_name", "titrant_conc", "theta_true"}
+
+
+def test_params_file_theta_gc_override_passed_to_thermo(mocker, tmp_path):
+ """theta_gc_override passed to thermo_to_growth contains measured genotypes."""
+ csv = _make_params_csv(
+ tmp_path,
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "wt,0.99,0.01,-4.1,2.0\n"
+ "A47V,0.97,0.03,-3.8,1.8\n",
+ )
+ config = _base_config_with_binding(csv)
+ mock_thermo, _ = _patch_for_params_file(mocker, config)
+
+ library_prediction(cf="config.yaml")
+
+ _, kwargs = mock_thermo.call_args
+ override = kwargs.get("theta_gc_override", {})
+ assert "wt" in override
+ assert "A47V" in override
+
+
+def test_unsupported_theta_component_raises(mocker, tmp_path):
+ """genotype_params_file with a non-Hill component must raise ValueError."""
+ csv = _make_params_csv(
+ tmp_path,
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\nwt,0.99,0.01,-4.1,2.0\n",
+ )
+ config = _base_config_with_binding(csv)
+ config["theta_component"] = "thermo.O2_C12_K5_U0_a.PK"
+
+ mocker.patch("tfscreen.util.read_yaml", return_value=config)
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.library_manager.LibraryManager"
+ ).return_value.build_library_df.return_value = pd.DataFrame({"genotype": ["wt"]})
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.build_sample_dataframes",
+ return_value=pd.DataFrame({"titrant_conc": [0.0]}),
+ )
+ mocker.patch("tfscreen.simulate.library_prediction.build_sim_data", return_value=MagicMock())
+ mocker.patch("tfscreen.simulate.library_prediction.jax.random.PRNGKey", return_value="k")
+
+ with pytest.raises(ValueError, match="hill"):
+ library_prediction(cf="config.yaml")
+
+
+def test_params_file_and_genotypes_coexist(mocker, tmp_path):
+ """genotype_params_file and genotypes can coexist; both appear in binding_df."""
+ csv = _make_params_csv(
+ tmp_path,
+ "genotype,theta_low,theta_high,log_hill_K,hill_n\n"
+ "A47V,0.97,0.03,-3.8,1.8\n",
+ )
+ config = _base_config_with_binding(csv, genotypes=["wt"])
+ mock_thermo, mock_sim_data = _patch_for_params_file(mocker, config)
+
+ # Patch sample_theta_prior (called for simulated WT in the 'genotypes' path)
+ import numpy as np
+ mock_wt_gc = np.array([[0.99, 0.50, 0.01]]) # shape (1, 3): 1 genotype × 3 concs
+ mock_theta_param = MagicMock()
+ mocker.patch(
+ "tfscreen.simulate.library_prediction.sample_theta_prior",
+ return_value=(mock_wt_gc, mock_theta_param),
+ )
+
+ _, _, _, _, binding_df = library_prediction(cf="config.yaml")
+
+ assert binding_df is not None
+ genos = set(binding_df["genotype"].unique())
+ # WT from simulated path + A47V from params file
+ assert "A47V" in genos
+ assert "wt" in genos
diff --git a/tests/tfscreen/simulate/test_run_simulation.py b/tests/tfscreen/simulate/test_run_simulation.py
index d96d8ed4..4ebf74c8 100644
--- a/tests/tfscreen/simulate/test_run_simulation.py
+++ b/tests/tfscreen/simulate/test_run_simulation.py
@@ -14,10 +14,11 @@ def mock_config():
def mock_result_dfs():
lib_df = pd.DataFrame({"lib": [1]})
pheno_df = pd.DataFrame({"pheno": [1]})
- ddG_df = pd.DataFrame({"ddG": [1]})
+ theta_df = pd.DataFrame({"theta": [1]})
+ params_df = pd.DataFrame({"genotype": ["wt"], "dk_geno": [0.0], "activity": [1.0]})
sample_df = pd.DataFrame({"sample": [1]})
counts_df = pd.DataFrame({"counts": [1]})
- return lib_df, pheno_df, ddG_df, sample_df, counts_df
+ return lib_df, pheno_df, theta_df, params_df, sample_df, counts_df
# -----------------------------------------------------------------------------
# Tests for _setup_file_output
@@ -67,17 +68,20 @@ def test_run_simulation_config_error():
run_simulation("nonexistent_config.yaml", None)
def test_run_simulation_success(mocker, mock_config, mock_result_dfs):
- lib_df, pheno_df, ddG_df, sample_df, counts_df = mock_result_dfs
-
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df = mock_result_dfs
+
# Mock dependencies
mocker.patch("tfscreen.util.read_yaml", return_value=mock_config)
- mock_lib_pred = mocker.patch("tfscreen.simulate.run_simulation.library_prediction", return_value=(lib_df, pheno_df, ddG_df))
+ mock_lib_pred = mocker.patch(
+ "tfscreen.simulate.run_simulation.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None),
+ )
mock_sel_exp = mocker.patch("tfscreen.simulate.run_simulation.selection_experiment", return_value=(sample_df, counts_df))
-
- # Mock file output setup
- mock_file_dict = {"library": "lib.csv", "phenotype": "pheno.csv", "genotype_theta": "theta.csv", "sample": "sample.csv", "counts": "counts.csv"}
+
+ # Mock file output setup — phenotype removed, parameters added
+ mock_file_dict = {"library": "lib.csv", "parameters": "params.csv", "genotype_theta": "theta.csv", "sample": "sample.csv", "counts": "counts.csv"}
mocker.patch("tfscreen.simulate.run_simulation._setup_file_output", return_value=mock_file_dict)
-
+
# Mock to_csv
mocker.patch.object(pd.DataFrame, "to_csv")
@@ -86,19 +90,22 @@ def test_run_simulation_success(mocker, mock_config, mock_result_dfs):
# Verify calls
mock_lib_pred.assert_called_once_with(mock_config)
mock_sel_exp.assert_called_once_with(mock_config, lib_df, pheno_df)
-
+
assert results["library"].equals(lib_df)
assert results["counts"].equals(counts_df)
-
- # Verify file writing
- assert pd.DataFrame.to_csv.call_count == 5
- # Logic for checking arguments could be added but call_count gives good confidence here given the simple loop logic
+ assert results["parameters"].equals(params_df)
+
+ # Verify file writing (5 files: library, parameters, genotype_theta, sample, counts)
+ assert pd.DataFrame.to_csv.call_count == 5
def test_run_simulation_no_output(mocker, mock_config, mock_result_dfs):
- lib_df, pheno_df, ddG_df, sample_df, counts_df = mock_result_dfs
-
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df = mock_result_dfs
+
mocker.patch("tfscreen.util.read_yaml", return_value=mock_config)
- mocker.patch("tfscreen.simulate.run_simulation.library_prediction", return_value=(lib_df, pheno_df, ddG_df))
+ mocker.patch(
+ "tfscreen.simulate.run_simulation.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None),
+ )
mocker.patch("tfscreen.simulate.run_simulation.selection_experiment", return_value=(sample_df, counts_df))
# Explicitly check _setup_file_output is called with None
diff --git a/tests/tfscreen/simulate/test_sample_theta.py b/tests/tfscreen/simulate/test_sample_theta.py
index 5b073c2c..a55c7605 100644
--- a/tests/tfscreen/simulate/test_sample_theta.py
+++ b/tests/tfscreen/simulate/test_sample_theta.py
@@ -1,8 +1,9 @@
import pytest
import numpy as np
+import pandas as pd
from unittest.mock import MagicMock, patch
-from tfscreen.simulate.sample_theta import sample_theta_prior, _EXCLUDED
+from tfscreen.simulate.sample_theta import sample_theta_prior, _EXCLUDED, _greedy_maximin, sample_theta_stratified
# ----------------------------------------------------------------------------
@@ -105,3 +106,249 @@ def test_define_model_called_with_correct_args(mock_sim_data):
sample_theta_prior("hill_geno", mock_sim_data, rng_key=0)
mock_module.define_model.assert_called_once_with("theta", mock_sim_data, mock_priors)
+
+
+# ----------------------------------------------------------------------------
+# Perturbation-path dispatch (simulate function present)
+# ----------------------------------------------------------------------------
+
+def _make_simulate_module(G, C):
+ """Return a mock theta module that exposes a real simulate() function."""
+ mock_module = MagicMock()
+ mock_module.get_sim_hyperparameters.return_value = {"alpha": 1.0}
+ mock_module.SimPriors.return_value = MagicMock()
+
+ expected_theta = np.ones((G, C)) * 0.7
+ expected_param = MagicMock(name="theta_param")
+
+ def fake_simulate(name, data, sim_priors, rng_key):
+ return expected_theta, expected_param
+
+ mock_module.simulate = fake_simulate
+ mock_module._expected_theta = expected_theta
+ mock_module._expected_param = expected_param
+ return mock_module
+
+
+def test_dispatch_uses_simulate_when_present(mock_sim_data):
+ G = mock_sim_data.num_genotype
+ C = mock_sim_data.num_titrant_conc
+ mock_module = _make_simulate_module(G, C)
+
+ with patch("tfscreen.simulate.sample_theta.model_registry",
+ {"theta": {"hill_geno": mock_module}}):
+ theta_gc, theta_param = sample_theta_prior("hill_geno", mock_sim_data, rng_key=0)
+
+ np.testing.assert_array_equal(theta_gc, mock_module._expected_theta)
+ assert theta_param is mock_module._expected_param
+ mock_module.define_model.assert_not_called()
+
+
+def test_dispatch_prior_path_when_no_simulate(mock_sim_data):
+ """Prior-predictive path is used when simulate is not a real function."""
+ G, C = mock_sim_data.num_genotype, mock_sim_data.num_titrant_conc
+ mock_module = _make_mock_module(G, C)
+ # MagicMock.simulate is a MagicMock, not a function — should fall through to prior path
+
+ with patch("tfscreen.simulate.sample_theta.model_registry",
+ {"theta": {"hill_geno": mock_module}}):
+ with patch("tfscreen.simulate.sample_theta.handlers"):
+ sample_theta_prior("hill_geno", mock_sim_data, rng_key=0)
+
+ mock_module.define_model.assert_called_once()
+
+
+def test_sim_priors_overrides_applied(mock_sim_data):
+ G = mock_sim_data.num_genotype
+ C = mock_sim_data.num_titrant_conc
+ mock_module = _make_simulate_module(G, C)
+
+ overrides = {"alpha": 9.9}
+ with patch("tfscreen.simulate.sample_theta.model_registry",
+ {"theta": {"hill_geno": mock_module}}):
+ sample_theta_prior("hill_geno", mock_sim_data, rng_key=0,
+ sim_priors_overrides=overrides)
+
+ call_kwargs = mock_module.SimPriors.call_args[1]
+ assert call_kwargs["alpha"] == 9.9
+
+
+# ----------------------------------------------------------------------------
+# _greedy_maximin
+# ----------------------------------------------------------------------------
+
+def test_greedy_maximin_returns_correct_count():
+ rng = np.random.default_rng(0)
+ theta_gc = rng.uniform(0, 1, (20, 5))
+ selected = _greedy_maximin(theta_gc, 4)
+ assert len(selected) == 4
+
+
+def test_greedy_maximin_no_duplicates():
+ rng = np.random.default_rng(1)
+ theta_gc = rng.uniform(0, 1, (50, 8))
+ selected = _greedy_maximin(theta_gc, 10)
+ assert len(selected) == len(set(selected.tolist()))
+
+
+def test_greedy_maximin_selects_spread_curves():
+ """The selected set should span the theta space broadly."""
+ # 50 curves spread across [0,1]; 2 selected should have large distance between them
+ rng = np.random.default_rng(7)
+ theta_gc = rng.uniform(0, 1, (50, 4))
+ selected = _greedy_maximin(theta_gc, 2)
+ a, b = theta_gc[selected[0]], theta_gc[selected[1]]
+ dist = np.sqrt(np.sum((a - b) ** 2))
+ # Two randomly chosen points from 50 uniform in [0,1]^4 would average ~0.87
+ # in max-pairwise distance; our greedy selection should comfortably exceed 0.5
+ assert dist > 0.5
+
+
+def test_greedy_maximin_all_when_n_equals_pool():
+ theta_gc = np.eye(5)
+ selected = _greedy_maximin(theta_gc, 5)
+ assert set(selected.tolist()) == set(range(5))
+
+
+def test_greedy_maximin_all_when_n_exceeds_pool():
+ theta_gc = np.eye(3)
+ selected = _greedy_maximin(theta_gc, 10)
+ assert list(selected) == [0, 1, 2]
+
+
+# ----------------------------------------------------------------------------
+# sample_theta_stratified
+# ----------------------------------------------------------------------------
+
+def _make_stratified_mock_module(n_binding_concs, n_growth_concs):
+ """
+ Return a mock theta module for sample_theta_stratified tests.
+
+ sample_theta_stratified now calls sample_theta_prior 2*pool_size times:
+ pool_size calls at binding concentrations, then pool_size calls at growth
+ concentrations, alternating binding/growth per pool member. Each call
+ returns a (1, C) array for the single-genotype sim_data.
+ """
+ mock_module = MagicMock()
+ mock_module.get_hyperparameters.return_value = {}
+ mock_module.ModelPriors.return_value = MagicMock()
+ mock_module.define_model.return_value = MagicMock(name="theta_param")
+
+ rng = np.random.default_rng(42)
+
+ def run_model_side_effect(theta_param, sim_data):
+ # Detect which sim_data is being used by its concentration count
+ n_conc = sim_data.num_titrant_conc
+ return rng.uniform(0, 1, (1, n_conc, 1))
+
+ mock_module.run_model.side_effect = run_model_side_effect
+ return mock_module
+
+
+def _make_stratified_sim_data(n_conc):
+ sd = MagicMock()
+ sd.num_genotype = 1
+ sd.num_titrant_conc = n_conc
+ return sd
+
+
+def test_stratified_output_shapes():
+ pool_size, n_select, n_binding_concs, n_growth_concs = 8, 3, 6, 3
+ mock_module = _make_stratified_mock_module(n_binding_concs, n_growth_concs)
+
+ binding_sample_df = pd.DataFrame({"titrant_conc": np.linspace(0, 1, n_binding_concs)})
+ growth_sample_df = pd.DataFrame({"titrant_conc": np.linspace(0, 1, n_growth_concs)})
+
+ def build_sim_data_side(library_df, sample_df, thermo_data=None, skip_pairs=False):
+ n_conc = len(sample_df["titrant_conc"].unique())
+ return _make_stratified_sim_data(n_conc)
+
+ with patch("tfscreen.simulate.sample_theta.model_registry",
+ {"theta": {"hill_geno": mock_module}}), \
+ patch("tfscreen.simulate.sample_theta.handlers"), \
+ patch("tfscreen.simulate.sample_theta.build_sim_data",
+ side_effect=build_sim_data_side):
+ binding_gc, growth_gc = sample_theta_stratified(
+ "hill_geno", binding_sample_df, growth_sample_df,
+ rng_key=0, n_select=n_select, pool_size=pool_size,
+ )
+
+ assert binding_gc.shape == (n_select, n_binding_concs)
+ assert growth_gc.shape == (n_select, n_growth_concs)
+
+
+def test_stratified_raises_when_n_select_exceeds_pool():
+ binding_sample_df = pd.DataFrame({"titrant_conc": [0.0, 1.0]})
+ growth_sample_df = pd.DataFrame({"titrant_conc": [0.5]})
+
+ with pytest.raises(ValueError, match="n_select"):
+ sample_theta_stratified(
+ "hill_geno", binding_sample_df, growth_sample_df,
+ rng_key=0, n_select=10, pool_size=5,
+ )
+
+
+def test_stratified_pool_size_calls():
+ """sample_theta_prior must be called 2*pool_size times (once per member per conc set)."""
+ pool_size, n_binding_concs, n_growth_concs = 5, 4, 2
+ mock_module = _make_stratified_mock_module(n_binding_concs, n_growth_concs)
+
+ binding_sample_df = pd.DataFrame({"titrant_conc": np.linspace(0, 1, n_binding_concs)})
+ growth_sample_df = pd.DataFrame({"titrant_conc": [0.0, 1.0]})
+
+ def build_sim_data_side(library_df, sample_df, thermo_data=None, skip_pairs=False):
+ return _make_stratified_sim_data(len(sample_df["titrant_conc"].unique()))
+
+ with patch("tfscreen.simulate.sample_theta.model_registry",
+ {"theta": {"hill_geno": mock_module}}), \
+ patch("tfscreen.simulate.sample_theta.handlers"), \
+ patch("tfscreen.simulate.sample_theta.build_sim_data",
+ side_effect=build_sim_data_side):
+ sample_theta_stratified(
+ "hill_geno", binding_sample_df, growth_sample_df,
+ rng_key=0, n_select=2, pool_size=pool_size,
+ )
+
+ assert mock_module.run_model.call_count == 2 * pool_size
+
+
+def test_stratified_ignores_simulate_path():
+ """sample_theta_stratified must use prior-predictive even when simulate() is present.
+
+ Mutation-decomposed components (e.g. hill_mut) with a wt-only pool library
+ have M=0 mutations and their simulate() path produces zero deltas for every
+ pool member. The prior-predictive path must be used instead so that WT-level
+ parameters are sampled independently per pool member.
+ """
+ pool_size, n_binding_concs, n_growth_concs = 4, 3, 2
+ mock_module = _make_stratified_mock_module(n_binding_concs, n_growth_concs)
+
+ # Add a real simulate() function — stratified path must NOT call it.
+ simulate_called = []
+
+ def fake_simulate(name, data, sim_priors, rng_key):
+ simulate_called.append(True)
+ return np.ones((data.num_genotype, data.num_titrant_conc)) * 0.5, MagicMock()
+
+ mock_module.simulate = fake_simulate
+ mock_module.get_sim_hyperparameters.return_value = {}
+ mock_module.SimPriors.return_value = MagicMock()
+
+ binding_sample_df = pd.DataFrame({"titrant_conc": np.linspace(0, 1, n_binding_concs)})
+ growth_sample_df = pd.DataFrame({"titrant_conc": [0.0, 1.0]})
+
+ def build_sim_data_side(library_df, sample_df, thermo_data=None, skip_pairs=False):
+ return _make_stratified_sim_data(len(sample_df["titrant_conc"].unique()))
+
+ with patch("tfscreen.simulate.sample_theta.model_registry",
+ {"theta": {"hill_geno": mock_module}}), \
+ patch("tfscreen.simulate.sample_theta.handlers"), \
+ patch("tfscreen.simulate.sample_theta.build_sim_data",
+ side_effect=build_sim_data_side):
+ sample_theta_stratified(
+ "hill_geno", binding_sample_df, growth_sample_df,
+ rng_key=0, n_select=2, pool_size=pool_size,
+ )
+
+ assert not simulate_called, "simulate() must not be called from sample_theta_stratified"
+ assert mock_module.run_model.call_count == 2 * pool_size
diff --git a/tests/tfscreen/simulate/test_selection_experiment.py b/tests/tfscreen/simulate/test_selection_experiment.py
index db1468d2..7915b92f 100644
--- a/tests/tfscreen/simulate/test_selection_experiment.py
+++ b/tests/tfscreen/simulate/test_selection_experiment.py
@@ -19,6 +19,7 @@
_sim_transform_and_mix,
_sim_growth,
MULTI_PLASMID_COMBINE_FCNS,
+ SIMULATE_KNOWN_KEYS,
_sim_sequencing,
_calc_genotype_cfu0,
_compute_kt,
@@ -42,8 +43,7 @@ def base_config() -> dict:
"lib_assembly_skew_sigma": 0.5,
"transformation_poisson_lambda": 0.8,
"tube_noise_sigma": 0.002,
- "final_cfu_pct_err": 0.03,
- "random_seed": 42,
+ "seed": 42,
"cfu0": 1.0e7,
"total_num_reads": 5_000_000,
"transform_sizes": {
@@ -198,7 +198,7 @@ def test_check_cf_loads_from_path(mocker, base_config: dict):
# Assert that the loader was called correctly
mock_loader.assert_called_once_with(dummy_path)
# Assert that the loaded config was processed
- assert validated_cf["random_seed"] == 42
+ assert validated_cf["seed"] == 42
@pytest.mark.parametrize("key_to_pop, nested_key_to_pop, bad_value, match_error", [
# Test required numerical key missing
@@ -1271,4 +1271,43 @@ def test_selection_experiment_end_to_end(mocker, base_config: dict,
assert counts_df.shape[0] == base_phenotype_df.shape[0]
total_reads = base_config["total_num_reads"]
- assert np.isclose(counts_df["counts"].sum(), total_reads, rtol=0.01)
\ No newline at end of file
+ assert np.isclose(counts_df["counts"].sum(), total_reads, rtol=0.01)
+
+
+# ----------------------------------------------------------------------------
+# test SIMULATE_KNOWN_KEYS / unknown-key validation in _check_cf
+# ----------------------------------------------------------------------------
+
+def test_simulate_known_keys_is_frozenset():
+ assert isinstance(SIMULATE_KNOWN_KEYS, frozenset)
+ assert len(SIMULATE_KNOWN_KEYS) > 0
+
+
+def test_check_cf_unknown_key_raises(base_config: dict):
+ bad = dict(base_config)
+ bad["not_a_real_key"] = 42
+ with pytest.raises(ValueError, match="not_a_real_key"):
+ _check_cf(bad)
+
+
+def test_check_cf_unknown_key_error_mentions_label(base_config: dict):
+ bad = dict(base_config)
+ bad["typo_key"] = "oops"
+ with pytest.raises(ValueError, match="simulate config"):
+ _check_cf(bad)
+
+
+def test_check_cf_multiple_unknown_keys(base_config: dict):
+ bad = dict(base_config)
+ bad["key_one"] = 1
+ bad["key_two"] = 2
+ with pytest.raises(ValueError) as exc_info:
+ _check_cf(bad)
+ msg = str(exc_info.value)
+ assert "key_one" in msg
+ assert "key_two" in msg
+
+
+def test_check_cf_all_known_keys_accepted(base_config: dict):
+ # Every key in the fixture must already be a known key; no error should be raised.
+ _check_cf(base_config)
\ No newline at end of file
diff --git a/tests/tfscreen/simulate/test_setup_sim_grid_cli.py b/tests/tfscreen/simulate/test_setup_sim_grid_cli.py
index 4fe0fe5b..4c44b570 100644
--- a/tests/tfscreen/simulate/test_setup_sim_grid_cli.py
+++ b/tests/tfscreen/simulate/test_setup_sim_grid_cli.py
@@ -24,7 +24,7 @@ def base_config(tmp_path):
"reading_frame": 0,
"observable_calculator": "lac",
"tube_noise_sigma": 0.01,
- "random_seed": None,
+ "seed": None,
}
p = tmp_path / "simulate_config.yaml"
p.write_text(yaml.dump(cfg))
@@ -151,8 +151,8 @@ def test_setup_sim_grid_cartesian_product(tmp_path, base_config):
- name: seed
variants:
- - random_seed: 0
- - random_seed: 42
+ - seed: 0
+ - seed: 42
"""
grid_path = tmp_path / "grid.yaml"
grid_path.write_text(grid)
diff --git a/tests/tfscreen/simulate/test_simulate_cli.py b/tests/tfscreen/simulate/test_simulate_cli.py
new file mode 100644
index 00000000..c448de1f
--- /dev/null
+++ b/tests/tfscreen/simulate/test_simulate_cli.py
@@ -0,0 +1,282 @@
+"""Tests for tfs-simulate CLI (run_simulation_from_config)."""
+
+import os
+import pytest
+from unittest.mock import patch, MagicMock
+import pandas as pd
+
+from tfscreen.simulate.scripts.simulate_cli import run_simulation_from_config
+from tfscreen.util.cli.generalized_main import generalized_main
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_mock_dfs():
+ lib_df = pd.DataFrame({"genotype": ["wt"], "sub_library": ["spiked"]})
+ pheno_df = pd.DataFrame({"genotype": ["wt"]})
+ theta_df = pd.DataFrame({"genotype": ["wt"]})
+ params_df = pd.DataFrame({"genotype": ["wt"], "dk_geno": [0.0], "activity": [1.0]})
+ sample_df = pd.DataFrame({"sample": [0]}, index=[0])
+ counts_df = pd.DataFrame({"sample": [0], "genotype": ["wt"], "counts": [1]})
+ growth_df = pd.DataFrame({"genotype": ["wt"]})
+ return lib_df, pheno_df, theta_df, params_df, sample_df, counts_df, growth_df
+
+
+# ---------------------------------------------------------------------------
+# --seed CLI argument
+# ---------------------------------------------------------------------------
+
+def test_seed_cli_parsed_as_int():
+ """--seed is registered with type=int so integer strings are accepted."""
+ captured = {}
+
+ def fake_run(config_file, output_dir, output_prefix="tfs_sim_",
+ num_replicates=2, seed=None):
+ captured["seed"] = seed
+
+ generalized_main(
+ fake_run,
+ argv=["config.yaml", "out_dir", "--seed", "42"],
+ manual_arg_types={"seed": int},
+ )
+ assert captured["seed"] == 42
+ assert isinstance(captured["seed"], int)
+
+
+def test_seed_cli_defaults_to_none():
+ """Omitting --seed leaves seed as None."""
+ captured = {}
+
+ def fake_run(config_file, output_dir, output_prefix="tfs_sim_",
+ num_replicates=2, seed=None):
+ captured["seed"] = seed
+
+ generalized_main(
+ fake_run,
+ argv=["config.yaml", "out_dir"],
+ manual_arg_types={"seed": int},
+ )
+ assert captured["seed"] is None
+
+
+# ---------------------------------------------------------------------------
+# seed override behaviour in run_simulation_from_config
+# ---------------------------------------------------------------------------
+
+@pytest.fixture()
+def patched_simulation(tmp_path):
+ """Patch all I/O so run_simulation_from_config can run without real data."""
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df, growth_df = _make_mock_dfs()
+
+ with patch("tfscreen.util.read_yaml", return_value={"seed": 99}) as mock_yaml, \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.selection_experiment",
+ return_value=(sample_df, counts_df)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.counts_to_lncfu",
+ return_value=growth_df), \
+ patch.object(pd.DataFrame, "to_csv"):
+ yield mock_yaml, tmp_path
+
+
+def test_seed_overrides_config(patched_simulation):
+ """When seed is given it replaces seed before seeding the RNG."""
+ mock_yaml, tmp_path = patched_simulation
+
+ with patch("tfscreen.simulate.scripts.simulate_cli.np.random.default_rng",
+ wraps=lambda s: MagicMock()) as mock_rng:
+ run_simulation_from_config("config.yaml", str(tmp_path), seed=7)
+
+ mock_rng.assert_called_once_with(7)
+
+
+def test_seed_none_preserves_config(patched_simulation):
+ """When seed=None the config's seed is left unchanged."""
+ mock_yaml, tmp_path = patched_simulation
+
+ with patch("tfscreen.simulate.scripts.simulate_cli.np.random.default_rng",
+ wraps=lambda s: MagicMock()) as mock_rng:
+ run_simulation_from_config("config.yaml", str(tmp_path), seed=None)
+
+ mock_rng.assert_called_once_with(99)
+
+
+# ---------------------------------------------------------------------------
+# Output file names
+# ---------------------------------------------------------------------------
+
+def test_writes_parameters_not_phenotype(tmp_path):
+ """run_simulation_from_config must write parameters.csv, not phenotype.csv."""
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df, growth_df = _make_mock_dfs()
+
+ written_paths = []
+
+ def capture_csv(self_df, path, **kwargs):
+ written_paths.append(str(path))
+
+ with patch("tfscreen.util.read_yaml", return_value={"seed": 0}), \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.selection_experiment",
+ return_value=(sample_df, counts_df)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.counts_to_lncfu",
+ return_value=growth_df), \
+ patch.object(pd.DataFrame, "to_csv", capture_csv):
+ run_simulation_from_config("config.yaml", str(tmp_path))
+
+ written = "\n".join(written_paths)
+ assert "parameters" in written, "parameters.csv must be written"
+ assert "phenotype" not in written, "phenotype.csv must NOT be written"
+
+
+def test_output_file_names_include_expected_stems(tmp_path):
+ """library, parameters, genotype_theta, and growth CSVs are all written."""
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df, growth_df = _make_mock_dfs()
+
+ written_paths = []
+
+ def capture_csv(self_df, path, **kwargs):
+ written_paths.append(str(path))
+
+ with patch("tfscreen.util.read_yaml", return_value={"seed": 0}), \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.selection_experiment",
+ return_value=(sample_df, counts_df)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.counts_to_lncfu",
+ return_value=growth_df), \
+ patch.object(pd.DataFrame, "to_csv", capture_csv):
+ run_simulation_from_config("config.yaml", str(tmp_path))
+
+ written = "\n".join(written_paths)
+ for stem in ("library", "parameters", "genotype_theta", "growth"):
+ assert stem in written, f"Expected '{stem}' CSV to be written"
+
+
+# ---------------------------------------------------------------------------
+# _generate_binding_data: uses pre-computed theta from library_prediction
+# ---------------------------------------------------------------------------
+
+def test_generate_binding_data_uses_precomputed_theta():
+ """_generate_binding_data must use theta_true from binding_theta_df,
+ not re-sample from the prior."""
+ import numpy as np
+ from tfscreen.simulate.scripts.simulate_cli import _generate_binding_data
+
+ binding_cfg = {
+ "genotypes": ["wt", "M1A"],
+ "titrant_name": "iptg",
+ "titrant_conc": [0.0, 1.0],
+ "noise": 0.0,
+ }
+ rng = np.random.default_rng(0)
+ binding_theta_df = pd.DataFrame([
+ {"genotype": "wt", "titrant_conc": 0.0, "theta_true": 0.1},
+ {"genotype": "wt", "titrant_conc": 1.0, "theta_true": 0.9},
+ {"genotype": "M1A", "titrant_conc": 0.0, "theta_true": 0.2},
+ {"genotype": "M1A", "titrant_conc": 1.0, "theta_true": 0.8},
+ ])
+
+ result = _generate_binding_data(binding_cfg, rng, binding_theta_df)
+
+ assert set(result.columns) >= {"genotype", "titrant_name", "titrant_conc",
+ "theta_obs", "theta_std"}
+ wt_row = result[(result["genotype"] == "wt") & (result["titrant_conc"] == 1.0)]
+ assert float(wt_row["theta_obs"].iloc[0]) == pytest.approx(0.9)
+
+
+def test_generate_binding_data_missing_genotype_raises():
+ """Raises ValueError when a genotype/conc pair is absent from binding_theta_df."""
+ import numpy as np
+ from tfscreen.simulate.scripts.simulate_cli import _generate_binding_data
+
+ binding_cfg = {
+ "genotypes": ["wt", "A2V"], # A2V is valid but not in binding_theta_df
+ "titrant_name": "iptg",
+ "titrant_conc": [1.0],
+ "noise": 0.0,
+ }
+ rng = np.random.default_rng(0)
+ binding_theta_df = pd.DataFrame([
+ {"genotype": "wt", "titrant_conc": 1.0, "theta_true": 0.5},
+ ])
+
+ with pytest.raises(ValueError, match="No pre-computed theta"):
+ _generate_binding_data(binding_cfg, rng, binding_theta_df)
+
+
+# ---------------------------------------------------------------------------
+# Rejecting the old theta_rng_seed key
+# ---------------------------------------------------------------------------
+
+def test_theta_rng_seed_rejected_as_unknown_key(tmp_path):
+ """A config containing theta_rng_seed must be rejected with an error
+ (it is no longer a recognized key)."""
+ from tfscreen.simulate.selection_experiment import _check_cf
+
+ cf = {
+ "theta_component": "hill_geno",
+ "theta_rng_seed": 0, # old key — must now be unknown
+ "condition_blocks": [],
+ "growth": {},
+ "transform_sizes": {},
+ "library_mixture": {},
+ "lib_assembly_skew_sigma": 0.0,
+ "transformation_poisson_lambda": 1,
+ "multi_plasmid_combine_fcn": "gmean",
+ "cfu0": 1e7,
+ "tube_noise_sigma": 0.0,
+ "total_num_reads": 1000,
+ "prob_index_hop": 0.0,
+ "seed": 0,
+ }
+ with pytest.raises(Exception): # check_unknown_keys raises ValueError
+ _check_cf(cf)
+
+
+# ---------------------------------------------------------------------------
+# input-config.yaml output
+# ---------------------------------------------------------------------------
+
+def test_writes_input_config_yaml(tmp_path):
+ """run_simulation_from_config must write an input-config.yaml using yaml.dump."""
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df, growth_df = _make_mock_dfs()
+
+ import tfscreen.simulate.scripts.simulate_cli as cli_mod
+
+ dumped = {}
+
+ def capture_dump(data, fh, **kwargs):
+ dumped["data"] = data
+
+ with patch("tfscreen.util.read_yaml", return_value={"seed": 5}), \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.selection_experiment",
+ return_value=(sample_df, counts_df)), \
+ patch("tfscreen.simulate.scripts.simulate_cli.counts_to_lncfu",
+ return_value=growth_df), \
+ patch.object(pd.DataFrame, "to_csv"), \
+ patch("tfscreen.simulate.scripts.simulate_cli.yaml.dump", capture_dump):
+ run_simulation_from_config("config.yaml", str(tmp_path), output_prefix="test_")
+
+ assert "data" in dumped, "yaml.dump was not called"
+ assert dumped["data"]["seed"] == 5
+
+
+def test_input_config_yaml_existence_check(tmp_path):
+ """If input-config.yaml already exists, FileExistsError is raised before any work."""
+ lib_df, pheno_df, theta_df, params_df, sample_df, counts_df, growth_df = _make_mock_dfs()
+
+ existing_yaml = tmp_path / "tfs_sim_input-config.yaml"
+ existing_yaml.write_text("seed: 0\n")
+
+ with patch("tfscreen.util.read_yaml", return_value={"seed": 0}), \
+ patch("tfscreen.simulate.scripts.simulate_cli.library_prediction",
+ return_value=(lib_df, pheno_df, theta_df, params_df, None)) as mock_lib:
+ with pytest.raises(FileExistsError, match="input-config.yaml"):
+ run_simulation_from_config("config.yaml", str(tmp_path))
+
+ mock_lib.assert_not_called()
diff --git a/tests/tfscreen/simulate/test_thermo_to_growth.py b/tests/tfscreen/simulate/test_thermo_to_growth.py
index b4d1b970..d4b3a56b 100644
--- a/tests/tfscreen/simulate/test_thermo_to_growth.py
+++ b/tests/tfscreen/simulate/test_thermo_to_growth.py
@@ -10,6 +10,7 @@
_apply_growth_params,
_sample_horseshoe_activity,
_sample_hierarchical_activity,
+ _theta_param_to_df,
_ACTIVITY_COMPONENTS,
thermo_to_growth,
_THETA_RESCALE,
@@ -142,6 +143,28 @@ def test_assign_dk_geno_distribution(rng):
assert (result <= 0.02).all()
+def test_assign_dk_geno_fixed_value_zero():
+ genotypes = ["wt", "A1B", "C2D", "A1B/C2D"]
+ result = _assign_dk_geno(genotypes, fixed_value=0.0)
+ assert isinstance(result, pd.Series)
+ assert set(result.index) == set(genotypes)
+ np.testing.assert_array_equal(result.values, 0.0)
+
+
+def test_assign_dk_geno_fixed_value_nonzero():
+ genotypes = ["wt", "A1B", "C2D"]
+ result = _assign_dk_geno(genotypes, fixed_value=-0.05)
+ np.testing.assert_array_equal(result.values, -0.05)
+
+
+def test_assign_dk_geno_fixed_value_skips_rng():
+ """fixed_value must not call the RNG — result is identical regardless of seed."""
+ genotypes = ["wt", "A1B", "C2D"]
+ r1 = _assign_dk_geno(genotypes, fixed_value=0.0, rng=np.random.default_rng(1))
+ r2 = _assign_dk_geno(genotypes, fixed_value=0.0, rng=np.random.default_rng(2))
+ np.testing.assert_array_equal(r1.values, r2.values)
+
+
# ----------------------------------------------------------------------------
# test _apply_growth_params
# ----------------------------------------------------------------------------
@@ -306,7 +329,7 @@ def test_thermo_to_growth_integration(
sim_data = _make_sim_data_mock(test_genotypes, concs)
_patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
- phenotype_df, genotype_theta_df = thermo_to_growth(
+ phenotype_df, genotype_theta_df, parameters_df = thermo_to_growth(
genotypes=test_genotypes,
sim_data=sim_data,
sample_df=test_sample_df,
@@ -319,6 +342,7 @@ def test_thermo_to_growth_integration(
assert isinstance(phenotype_df, pd.DataFrame)
assert phenotype_df.shape[0] == 6
assert "theta" in phenotype_df.columns
+ # phenotype_df retains dk_geno and activity for selection_experiment
assert "dk_geno" in phenotype_df.columns
assert "activity" in phenotype_df.columns
assert "k_pre" in phenotype_df.columns
@@ -343,11 +367,63 @@ def test_thermo_to_growth_integration(
phenotype_df["k_sel"].to_numpy(), b_sel + activity * m_sel * theta + dk
)
- # genotype_theta_df: one row per unique genotype, columns theta_at_*mM
+ # genotype_theta_df: long form — one row per (genotype, titrant_name, titrant_conc)
assert isinstance(genotype_theta_df, pd.DataFrame)
- assert "genotype" in genotype_theta_df.columns
- theta_cols = [c for c in genotype_theta_df.columns if c.startswith("theta_at_")]
- assert len(theta_cols) == 2 # two unique concentrations
+ assert list(genotype_theta_df.columns) == ["genotype", "titrant_name",
+ "titrant_conc", "theta"]
+ # 3 unique genotypes × 2 unique concentrations = 6 rows
+ assert len(genotype_theta_df) == 6
+ # Each (genotype, titrant_name, titrant_conc) triple is unique
+ assert genotype_theta_df.duplicated(
+ subset=["genotype", "titrant_name", "titrant_conc"]
+ ).sum() == 0
+
+ # parameters_df: one row per unique genotype, dk_geno + activity always present
+ assert isinstance(parameters_df, pd.DataFrame)
+ assert parameters_df.shape[0] == 3 # 3 unique genotypes
+ assert list(parameters_df.columns[:3]) == ["genotype", "dk_geno", "activity"]
+ # wt always gets dk_geno == 0
+ wt_row = parameters_df[parameters_df["genotype"] == "wt"]
+ assert np.isclose(float(wt_row["dk_geno"].iloc[0]), 0.0)
+ # With activity_mut_scale=0, all activities == 1.0
+ np.testing.assert_allclose(parameters_df["activity"].values, 1.0)
+
+
+def test_genotype_theta_df_no_duplicates_with_repeated_genotypes(
+ mocker, test_sample_df, simple_growth_params
+):
+ """genotype_theta_df must have one row per unique genotype even when the
+ same genotype appears multiple times in the library (e.g. two sub-libraries)."""
+ # Supply 5 rows where "wt" and "A1B" each appear twice
+ genotypes_with_dups = ["wt", "A1B", "wt", "A1B", "A1B/C2D"]
+ concs = np.array([10.0, 100.0])
+ # theta_gc has 5 rows (one per library entry); duplicate genotypes get
+ # distinct values, but only one row per unique genotype should be kept.
+ theta_gc = np.array([
+ [0.1, 0.9], # wt (first occurrence)
+ [0.3, 0.7], # A1B (first occurrence)
+ [0.2, 0.8], # wt (second occurrence — should be dropped)
+ [0.4, 0.6], # A1B (second occurrence — should be dropped)
+ [0.5, 0.5], # A1B/C2D
+ ])
+ sim_data = _make_sim_data_mock(genotypes_with_dups, concs)
+ _patch_thermo_deps(mocker, genotypes_with_dups, test_sample_df, theta_gc)
+
+ _, genotype_theta_df, _ = thermo_to_growth(
+ genotypes=genotypes_with_dups,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ )
+
+ # 3 unique genotypes × 2 concentrations = 6 rows; no duplicated (geno, conc) pairs
+ assert genotype_theta_df["genotype"].nunique() == 3
+ assert len(genotype_theta_df) == 6
+ assert genotype_theta_df.duplicated(
+ subset=["genotype", "titrant_name", "titrant_conc"]
+ ).sum() == 0
def test_thermo_to_growth_propagates_rng(
@@ -362,7 +438,7 @@ def test_thermo_to_growth_propagates_rng(
mock_assign_dk = mocker.patch(
"tfscreen.simulate.thermo_to_growth._assign_dk_geno",
- return_value=pd.Series(0.0, index=["wt", "A1B", "A1B/C2D"]),
+ return_value=pd.Series({"wt": 0.0, "A1B": 0.0, "A1B/C2D": 0.0}),
)
thermo_to_growth(
@@ -380,6 +456,357 @@ def test_thermo_to_growth_propagates_rng(
assert passed_rng is rng
+def test_thermo_to_growth_dk_geno_zero_sets_all_to_zero(
+ mocker, test_genotypes, test_sample_df, simple_growth_params
+):
+ """dk_geno_zero=True must set dk_geno=0 for every genotype."""
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.5)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ phenotype_df, _, parameters_df = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ dk_geno_zero=True,
+ )
+
+ np.testing.assert_array_equal(phenotype_df["dk_geno"].values, 0.0)
+ np.testing.assert_array_equal(parameters_df["dk_geno"].values, 0.0)
+
+
+def test_thermo_to_growth_dk_geno_zero_passes_fixed_value(
+ mocker, test_genotypes, test_sample_df, simple_growth_params
+):
+ """dk_geno_zero=True must pass fixed_value=0.0 to _assign_dk_geno."""
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.zeros((3, 2))
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ mock_assign_dk = mocker.patch(
+ "tfscreen.simulate.thermo_to_growth._assign_dk_geno",
+ return_value=pd.Series({"wt": 0.0, "A1B": 0.0, "A1B/C2D": 0.0}),
+ )
+
+ thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ dk_geno_zero=True,
+ )
+
+ assert mock_assign_dk.call_args.kwargs.get("fixed_value") == 0.0
+
+
+def test_thermo_to_growth_dk_geno_zero_false_passes_none(
+ mocker, test_genotypes, test_sample_df, simple_growth_params
+):
+ """dk_geno_zero=False (default) must pass fixed_value=None to _assign_dk_geno."""
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.zeros((3, 2))
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ mock_assign_dk = mocker.patch(
+ "tfscreen.simulate.thermo_to_growth._assign_dk_geno",
+ return_value=pd.Series({"wt": 0.0, "A1B": 0.0, "A1B/C2D": 0.0}),
+ )
+
+ thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ dk_geno_zero=False,
+ )
+
+ assert mock_assign_dk.call_args.kwargs.get("fixed_value") is None
+
+
+# ============================================================================
+# test thermo_to_growth — theta_gc_override
+# ============================================================================
+
+class TestThetaGcOverride:
+ """theta_gc_override must replace theta rows for the named genotypes."""
+
+ def test_overridden_genotype_has_new_theta(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.zeros((3, 2)) # all-zero baseline
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ override_values = np.array([0.8, 0.9])
+ _, genotype_theta_df, _ = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_gc_override={"A1B": override_values},
+ )
+
+ a1b_rows = genotype_theta_df[genotype_theta_df["genotype"] == "A1B"]
+ a1b_rows = a1b_rows.sort_values("titrant_conc").reset_index(drop=True)
+ np.testing.assert_allclose(a1b_rows["theta"].values, override_values)
+
+ def test_non_overridden_genotypes_unchanged(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.zeros((3, 2))
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ _, genotype_theta_df, _ = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_gc_override={"A1B": np.array([0.8, 0.9])},
+ )
+
+ for geno in ("wt", "A1B/C2D"):
+ rows = genotype_theta_df[genotype_theta_df["genotype"] == geno]
+ np.testing.assert_array_equal(rows["theta"].values, 0.0)
+
+ def test_none_override_is_noop(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.3)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ _, genotype_theta_df, _ = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_gc_override=None,
+ )
+
+ np.testing.assert_array_equal(genotype_theta_df["theta"].values, 0.3)
+
+ def test_empty_override_is_noop(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.3)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ _, genotype_theta_df, _ = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_gc_override={},
+ )
+
+ np.testing.assert_array_equal(genotype_theta_df["theta"].values, 0.3)
+
+ def test_unknown_genotype_in_override_is_ignored(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.zeros((3, 2))
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+ _patch_thermo_deps(mocker, test_genotypes, test_sample_df, theta_gc)
+
+ _, genotype_theta_df, _ = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_gc_override={"NOTINGENO": np.array([0.99, 0.99])},
+ )
+
+ np.testing.assert_array_equal(genotype_theta_df["theta"].values, 0.0)
+
+
+# ============================================================================
+# test thermo_to_growth — theta_params_override
+# ============================================================================
+
+def _patch_thermo_deps_with_param(mocker, theta_gc, theta_param):
+ """Like _patch_thermo_deps but accepts a custom theta_param object."""
+ mocker.patch(
+ "tfscreen.simulate.thermo_to_growth.sample_theta_prior",
+ return_value=(theta_gc, theta_param),
+ )
+ mocker.patch(
+ "tfscreen.simulate.thermo_to_growth.set_categorical_genotype",
+ side_effect=lambda df: df,
+ )
+
+
+class TestThetaParamsOverride:
+ """theta_params_override must update Hill columns in parameters_df."""
+
+ def _make_hill_param(self, n_geno):
+ """Return a _MockThetaParam2D with shape (1, n_geno) and known values."""
+ rng = np.random.default_rng(5)
+ return _MockThetaParam2D(
+ theta_low=rng.random((1, n_geno)),
+ theta_high=rng.random((1, n_geno)),
+ log_hill_K=rng.standard_normal((1, n_geno)),
+ hill_n=np.abs(rng.standard_normal((1, n_geno))) + 0.5,
+ )
+
+ def test_overridden_genotype_columns_updated(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.5)
+ theta_param = self._make_hill_param(3)
+ _patch_thermo_deps_with_param(mocker, theta_gc, theta_param)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+
+ override = {"A1B": {"theta_low": 0.99, "theta_high": 0.01,
+ "log_hill_K": -4.0, "hill_n": 2.0}}
+ _, _, parameters_df = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_params_override=override,
+ )
+
+ row = parameters_df[parameters_df["genotype"] == "A1B"].iloc[0]
+ assert np.isclose(row["theta_low"], 0.99)
+ assert np.isclose(row["theta_high"], 0.01)
+ assert np.isclose(row["log_hill_K"], -4.0)
+ assert np.isclose(row["hill_n"], 2.0)
+
+ def test_non_overridden_genotypes_retain_original_values(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.5)
+ theta_param = self._make_hill_param(3)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+
+ # First: baseline call with no override to get reference values
+ _patch_thermo_deps_with_param(mocker, theta_gc, theta_param)
+ _, _, baseline_df = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ )
+ wt_theta_low_baseline = float(
+ baseline_df[baseline_df["genotype"] == "wt"]["theta_low"].iloc[0]
+ )
+ double_theta_low_baseline = float(
+ baseline_df[baseline_df["genotype"] == "A1B/C2D"]["theta_low"].iloc[0]
+ )
+
+ # Second: override call for A1B only
+ _patch_thermo_deps_with_param(mocker, theta_gc, theta_param)
+ _, _, parameters_df = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_params_override={"A1B": {"theta_low": 0.99}},
+ )
+
+ wt_row = parameters_df[parameters_df["genotype"] == "wt"].iloc[0]
+ assert np.isclose(wt_row["theta_low"], wt_theta_low_baseline)
+ double_row = parameters_df[parameters_df["genotype"] == "A1B/C2D"].iloc[0]
+ assert np.isclose(double_row["theta_low"], double_theta_low_baseline)
+
+ def test_none_override_is_noop(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.5)
+ theta_param = self._make_hill_param(3)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+
+ _patch_thermo_deps_with_param(mocker, theta_gc, theta_param)
+ _, _, parameters_df_no_override = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ rng=np.random.default_rng(0),
+ theta_params_override=None,
+ )
+
+ _patch_thermo_deps_with_param(mocker, theta_gc, theta_param)
+ _, _, parameters_df_empty = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ rng=np.random.default_rng(0),
+ theta_params_override={},
+ )
+
+ pd.testing.assert_frame_equal(
+ parameters_df_no_override.reset_index(drop=True),
+ parameters_df_empty.reset_index(drop=True),
+ )
+
+ def test_unknown_column_in_override_is_silently_skipped(
+ self, mocker, test_genotypes, test_sample_df, simple_growth_params
+ ):
+ concs = np.array([10.0, 100.0])
+ theta_gc = np.full((3, 2), 0.5)
+ theta_param = self._make_hill_param(3)
+ _patch_thermo_deps_with_param(mocker, theta_gc, theta_param)
+ sim_data = _make_sim_data_mock(test_genotypes, concs)
+
+ # "nonexistent_col" is not in parameters_df; must not raise
+ _, _, parameters_df = thermo_to_growth(
+ genotypes=test_genotypes,
+ sim_data=sim_data,
+ sample_df=test_sample_df,
+ theta_component="mock",
+ theta_rng_key=0,
+ growth_params=simple_growth_params,
+ theta_params_override={"A1B": {"theta_low": 0.99,
+ "nonexistent_col": 42.0}},
+ )
+
+ assert "nonexistent_col" not in parameters_df.columns
+ row = parameters_df[parameters_df["genotype"] == "A1B"].iloc[0]
+ assert np.isclose(row["theta_low"], 0.99)
+
+
# ============================================================================
# test _sample_horseshoe_activity
# ============================================================================
@@ -713,13 +1140,17 @@ def test_horseshoe_activities_appear_in_phenotype_df(
"tfscreen.simulate.thermo_to_growth._sample_horseshoe_activity",
return_value=pd.Series(activity_map),
)
- phenotype_df, _ = thermo_to_growth(
+ phenotype_df, _, parameters_df = thermo_to_growth(
genotypes=test_genotypes, sim_data=sim_data,
activity_component="horseshoe_geno", **base_call_kwargs,
)
for geno, expected in activity_map.items():
rows = phenotype_df[phenotype_df["genotype"] == geno]
np.testing.assert_allclose(rows["activity"].values, expected)
+ # Same activity values must appear in parameters_df (one row per genotype)
+ for geno, expected in activity_map.items():
+ row = parameters_df[parameters_df["genotype"] == geno]
+ np.testing.assert_allclose(float(row["activity"].iloc[0]), expected)
# ============================================================================
@@ -751,7 +1182,7 @@ def test_zero_sigma_leaves_theta_unchanged(
"""theta_noise_sigma_logit=0 must not alter theta values."""
concs = np.array([10.0, 100.0])
theta_gc = np.array([[0.2, 0.8], [0.3, 0.7], [0.5, 0.5]])
- phenotype_df, _ = self._run(
+ phenotype_df, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=0.0,
rng=np.random.default_rng(0),
@@ -767,7 +1198,7 @@ def test_nonzero_sigma_perturbs_theta(
"""theta_noise_sigma_logit > 0 must produce theta values different from input."""
concs = np.array([10.0, 100.0])
theta_gc = np.full((3, 2), 0.5)
- phenotype_df, _ = self._run(
+ phenotype_df, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=1.0,
rng=np.random.default_rng(42),
@@ -781,7 +1212,7 @@ def test_noisy_theta_stays_in_unit_interval(
"""Even with large sigma_logit, theta must remain in (0, 1)."""
concs = np.array([10.0, 100.0])
theta_gc = np.array([[0.01, 0.99], [0.5, 0.5], [0.3, 0.7]])
- phenotype_df, _ = self._run(
+ phenotype_df, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=5.0,
rng=np.random.default_rng(0),
@@ -798,13 +1229,13 @@ def test_reproducible_with_same_rng(
concs = np.array([10.0, 100.0])
theta_gc = np.full((3, 2), 0.5)
- df1, _ = self._run(
+ df1, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=0.5,
rng=np.random.default_rng(7),
**base_kwargs,
)
- df2, _ = self._run(
+ df2, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=0.5,
rng=np.random.default_rng(7),
@@ -818,13 +1249,13 @@ def test_different_seeds_give_different_theta(
concs = np.array([10.0, 100.0])
theta_gc = np.full((3, 2), 0.5)
- df1, _ = self._run(
+ df1, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=0.5,
rng=np.random.default_rng(1),
**base_kwargs,
)
- df2, _ = self._run(
+ df2, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_noise_sigma_logit=0.5,
rng=np.random.default_rng(2),
@@ -902,12 +1333,12 @@ def test_passthrough_matches_default(
concs = np.array([10.0, 100.0])
theta_gc = np.array([[0.2, 0.8], [0.3, 0.7], [0.5, 0.5]])
- df_default, _ = self._run(
+ df_default, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
rng=np.random.default_rng(42),
**base_kwargs,
)
- df_passthrough, _ = self._run(
+ df_passthrough, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_rescale="passthrough",
rng=np.random.default_rng(42),
@@ -929,12 +1360,12 @@ def test_logit_changes_k_relative_to_passthrough(
concs = np.array([10.0, 100.0])
theta_gc = np.array([[0.2, 0.8], [0.3, 0.7], [0.5, 0.5]])
- df_pass, _ = self._run(
+ df_pass, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_rescale="passthrough",
**base_kwargs,
)
- df_logit, _ = self._run(
+ df_logit, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_rescale="logit",
**base_kwargs,
@@ -948,7 +1379,7 @@ def test_logit_theta_column_stays_in_unit_interval(
concs = np.array([10.0, 100.0])
theta_gc = np.array([[0.2, 0.8], [0.3, 0.7], [0.5, 0.5]])
- df, _ = self._run(
+ df, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_rescale="logit",
**base_kwargs,
@@ -967,7 +1398,7 @@ def test_logit_k_matches_hand_calculation(
# activity_mut_scale so only the logit transform matters.
theta_gc = np.full((3, 2), 0.5)
- df, _ = self._run(
+ df, _, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_rescale="logit",
activity_wt=1.0,
@@ -988,7 +1419,7 @@ def test_genotype_theta_df_stays_in_unit_interval_under_logit(
concs = np.array([10.0, 100.0])
theta_gc = np.array([[0.2, 0.8], [0.3, 0.7], [0.5, 0.5]])
- _, geno_theta_df = self._run(
+ _, geno_theta_df, _ = self._run(
mocker, test_genotypes, concs, theta_gc,
theta_rescale="logit",
**base_kwargs,
@@ -997,3 +1428,145 @@ def test_genotype_theta_df_stays_in_unit_interval_under_logit(
vals = geno_theta_df[theta_cols].values
assert np.all(vals > 0.0)
assert np.all(vals < 1.0)
+
+
+# ============================================================================
+# test _theta_param_to_df
+# ============================================================================
+
+from dataclasses import dataclass as _stdlib_dataclass
+
+
+@_stdlib_dataclass
+class _MockThetaParam2D:
+ """Minimal stand-in for a real ThetaParam with only 2-D per-genotype fields."""
+ theta_low: object
+ theta_high: object
+ log_hill_K: object
+ hill_n: object
+
+
+@_stdlib_dataclass
+class _MockThetaParamMixed:
+ """ThetaParam with 2-D per-genotype fields AND a 3-D population field."""
+ theta_low: object
+ mu: object # 3-D — should be ignored
+
+
+class TestThetaParamToDf:
+ """Unit tests for _theta_param_to_df."""
+
+ def _make_param_single_titrant(self, G):
+ """Return a _MockThetaParam2D with shape (1, G) fields."""
+ rng = np.random.default_rng(0)
+ return _MockThetaParam2D(
+ theta_low=rng.random((1, G)),
+ theta_high=rng.random((1, G)),
+ log_hill_K=rng.standard_normal((1, G)),
+ hill_n=np.abs(rng.standard_normal((1, G))),
+ )
+
+ def test_returns_dataframe(self):
+ genotypes = ["wt", "A1V", "B2G"]
+ sim_idx = np.array([0, 1, 2])
+ param = self._make_param_single_titrant(3)
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ assert isinstance(df, pd.DataFrame)
+
+ def test_genotype_column_present_and_correct(self):
+ genotypes = ["wt", "A1V", "B2G"]
+ sim_idx = np.array([0, 1, 2])
+ param = self._make_param_single_titrant(3)
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ assert "genotype" in df.columns
+ assert list(df["genotype"]) == genotypes
+
+ def test_genotype_is_first_column(self):
+ genotypes = ["wt", "A1V"]
+ sim_idx = np.array([0, 1])
+ param = self._make_param_single_titrant(2)
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ assert df.columns[0] == "genotype"
+
+ def test_single_titrant_fields_extracted(self):
+ """All four 2-D fields must appear as columns (T=1 is squeezed)."""
+ genotypes = ["wt", "A1V", "B2G"]
+ sim_idx = np.array([0, 1, 2])
+ param = self._make_param_single_titrant(3)
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ for fname in ["theta_low", "theta_high", "log_hill_K", "hill_n"]:
+ assert fname in df.columns, f"Missing column: {fname}"
+
+ def test_3d_field_is_skipped(self):
+ """A field with ndim=3 (e.g. population moments) must not appear."""
+ genotypes = ["wt", "A1V"]
+ sim_idx = np.array([0, 1])
+ param = _MockThetaParamMixed(
+ theta_low=np.array([[0.1, 0.2]]), # (1, 2) — kept
+ mu=np.array([[[0.5], [0.6]]]), # (1, 2, 1) — skipped
+ )
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ assert "theta_low" in df.columns
+ assert "mu" not in df.columns
+
+ def test_values_match_theta_param(self):
+ """Extracted values must match the raw theta_param arrays."""
+ genotypes = ["wt", "A1V", "B2G"]
+ sim_idx = np.array([0, 1, 2])
+ param = self._make_param_single_titrant(3)
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ np.testing.assert_allclose(df["theta_low"].values, param.theta_low[0])
+ np.testing.assert_allclose(df["theta_high"].values, param.theta_high[0])
+
+ def test_sim_indices_select_correct_rows(self):
+ """sim_indices must select the right columns from theta_param."""
+ # sim_data has 5 genotypes; we want 3 of them in reverse order
+ G_sim = 5
+ G_out = 3
+ sim_idx = np.array([4, 2, 0]) # select in reverse from sim_data
+ rng = np.random.default_rng(7)
+ theta_low_full = rng.random((1, G_sim))
+ param = _MockThetaParam2D(
+ theta_low=theta_low_full,
+ theta_high=rng.random((1, G_sim)),
+ log_hill_K=rng.standard_normal((1, G_sim)),
+ hill_n=np.abs(rng.standard_normal((1, G_sim))),
+ )
+ genotypes = ["g4", "g2", "g0"]
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ expected = theta_low_full[0, sim_idx]
+ np.testing.assert_allclose(df["theta_low"].values, expected)
+
+ def test_multi_titrant_creates_suffixed_columns(self):
+ """When T > 1, columns are named field_T0, field_T1, …"""
+ T, G = 2, 3
+ rng = np.random.default_rng(1)
+ param = _MockThetaParam2D(
+ theta_low=rng.random((T, G)),
+ theta_high=rng.random((T, G)),
+ log_hill_K=rng.standard_normal((T, G)),
+ hill_n=np.abs(rng.standard_normal((T, G))),
+ )
+ genotypes = ["wt", "A1V", "B2G"]
+ sim_idx = np.array([0, 1, 2])
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ assert "theta_low_T0" in df.columns
+ assert "theta_low_T1" in df.columns
+ assert "theta_low" not in df.columns # unsuffixed form absent
+
+ def test_non_dataclass_returns_genotype_only(self):
+ """A non-dataclass theta_param (e.g. MagicMock) yields only 'genotype' column."""
+ from unittest.mock import MagicMock
+ mock_param = MagicMock()
+ genotypes = ["wt", "A1V"]
+ sim_idx = np.array([0, 1])
+ df = _theta_param_to_df(mock_param, genotypes, sim_idx)
+ assert set(df.columns) == {"genotype"}
+ assert list(df["genotype"]) == genotypes
+
+ def test_one_row_per_genotype(self):
+ genotypes = ["wt", "A1V", "B2G", "C3H"]
+ sim_idx = np.array([0, 1, 2, 3])
+ param = self._make_param_single_titrant(4)
+ df = _theta_param_to_df(param, genotypes, sim_idx)
+ assert len(df) == 4
diff --git a/tests/tfscreen/tfmodel/components/growth/__init__.py b/tests/tfscreen/tfmodel/analysis/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth/__init__.py
rename to tests/tfscreen/tfmodel/analysis/__init__.py
diff --git a/tests/tfscreen/tfmodel/analysis/test_error_calibration.py b/tests/tfscreen/tfmodel/analysis/test_error_calibration.py
new file mode 100644
index 00000000..2da9b5bc
--- /dev/null
+++ b/tests/tfscreen/tfmodel/analysis/test_error_calibration.py
@@ -0,0 +1,723 @@
+"""
+Tests for tfscreen.tfmodel.analysis.error_calibration.
+
+Covers all public functions and the private SBC helpers that are tested by
+the now-deleted test_sbc.py.
+"""
+
+import os
+
+import h5py
+import numpy as np
+import pandas as pd
+import pytest
+
+from tfscreen.tfmodel.analysis.error_calibration import (
+ # Core
+ pit_from_samples,
+ pit_from_quantiles,
+ # Stats
+ calibration_curve,
+ pit_uniformity_test,
+ # Plots
+ plot_pit_histogram,
+ plot_calibration_curve,
+ # Single-run
+ calibration_summary,
+ # SBC
+ _find_pairs,
+ _load_h5_params,
+ compute_sbc_ranks,
+ summarize_sbc,
+)
+
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+def _write_h5(path, arrays):
+ """Write a dict of numpy arrays to an HDF5 file."""
+ with h5py.File(path, "w") as hf:
+ for key, val in arrays.items():
+ hf.create_dataset(key, data=val)
+
+
+# ============================================================
+# Part 1 — pit_from_samples
+# ============================================================
+
+class TestPitFromSamples:
+
+ def test_rank_zero_when_gt_below_all_posterior(self):
+ # true = -100; all posterior > -100 → PIT = 0
+ pit = pit_from_samples(np.array([-100.0]), np.ones((50, 1)))
+ np.testing.assert_allclose(pit, [0.0])
+
+ def test_rank_one_when_gt_above_all_posterior(self):
+ # true = 100; all posterior < 100 → PIT = 1
+ pit = pit_from_samples(np.array([100.0]), -np.ones((50, 1)))
+ np.testing.assert_allclose(pit, [1.0])
+
+ def test_rank_half_at_median(self):
+ # posterior is -10..10 (21 values), true = 0 → 10/21
+ post = np.arange(-10, 11, dtype=float).reshape(21, 1)
+ pit = pit_from_samples(np.array([0.0]), post)
+ np.testing.assert_allclose(pit, [10 / 21])
+
+ def test_nan_true_value_gives_nan(self):
+ post = np.ones((20, 1))
+ pit = pit_from_samples(np.array([np.nan]), post)
+ assert np.isnan(pit[0])
+
+ def test_mixed_nan_and_valid(self):
+ true_vals = np.array([0.0, np.nan, 1000.0])
+ post = np.zeros((20, 3)) # all samples = 0
+ pit = pit_from_samples(true_vals, post)
+ # true=0: mean(0 < 0) = 0
+ np.testing.assert_allclose(pit[0], 0.0)
+ assert np.isnan(pit[1])
+ # true=1000: mean(0 < 1000) = 1
+ np.testing.assert_allclose(pit[2], 1.0)
+
+ def test_multidimensional_true_vals(self):
+ rng = np.random.default_rng(0)
+ post = rng.normal(size=(100, 4))
+ true_vals = np.zeros(4)
+ pit = pit_from_samples(true_vals, post)
+ assert pit.shape == (4,)
+ assert np.all((pit >= 0) & (pit <= 1))
+
+ def test_output_in_unit_interval(self):
+ rng = np.random.default_rng(42)
+ post = rng.normal(size=(200, 10))
+ true_vals = rng.normal(size=10)
+ pit = pit_from_samples(true_vals, post)
+ assert np.all((pit >= 0) & (pit <= 1))
+
+ def test_scalar_true_val(self):
+ # true_vals can be a scalar if posterior is shape (S,)
+ post = np.linspace(-1, 1, 101)
+ pit = pit_from_samples(0.0, post)
+ # 50 values < 0 → 50/101
+ np.testing.assert_allclose(pit, 50 / 101)
+
+
+# ============================================================
+# Part 1 — pit_from_quantiles
+# ============================================================
+
+class TestPitFromQuantiles:
+
+ def _make_gaussian_quantiles(self, n=5):
+ """Return quantile_matrix and levels for N(0,1) at 9 levels."""
+ from scipy.stats import norm
+ levels = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
+ qmat = np.tile(norm.ppf(levels), (n, 1)) # same distribution for all n
+ return qmat, levels
+
+ def test_true_at_stored_quantile_returns_that_level(self):
+ from scipy.stats import norm
+ levels = np.array([0.1, 0.5, 0.9])
+ qmat = norm.ppf(levels).reshape(1, 3)
+ # True value exactly at the 50th percentile
+ pit = pit_from_quantiles(np.array([0.0]), qmat, levels)
+ np.testing.assert_allclose(pit, [0.5])
+
+ def test_true_below_all_quantiles_returns_zero(self):
+ levels = np.array([0.1, 0.5, 0.9])
+ qmat = np.array([[1.0, 2.0, 3.0]])
+ pit = pit_from_quantiles(np.array([-999.0]), qmat, levels)
+ np.testing.assert_allclose(pit, [0.0])
+
+ def test_true_above_all_quantiles_returns_one(self):
+ levels = np.array([0.1, 0.5, 0.9])
+ qmat = np.array([[1.0, 2.0, 3.0]])
+ pit = pit_from_quantiles(np.array([999.0]), qmat, levels)
+ np.testing.assert_allclose(pit, [1.0])
+
+ def test_linear_interpolation_between_levels(self):
+ # quantile values 0,1,2,3,4 at levels 0,0.25,0.50,0.75,1.0
+ levels = np.array([0.0, 0.25, 0.50, 0.75, 1.0])
+ qmat = np.array([[0.0, 1.0, 2.0, 3.0, 4.0]])
+ # True = 1.5: halfway between q=1.0 (level 0.25) and q=2.0 (level 0.5)
+ pit = pit_from_quantiles(np.array([1.5]), qmat, levels)
+ np.testing.assert_allclose(pit, [0.375], atol=1e-10)
+
+ def test_nan_true_value_gives_nan(self):
+ levels = np.array([0.1, 0.5, 0.9])
+ qmat = np.array([[1.0, 2.0, 3.0]])
+ pit = pit_from_quantiles(np.array([np.nan]), qmat, levels)
+ assert np.isnan(pit[0])
+
+ def test_nan_in_quantile_row_gives_nan(self):
+ levels = np.array([0.1, 0.5, 0.9])
+ qmat = np.array([[1.0, np.nan, 3.0]])
+ pit = pit_from_quantiles(np.array([2.0]), qmat, levels)
+ assert np.isnan(pit[0])
+
+ def test_shape_preserved(self):
+ qmat, levels = self._make_gaussian_quantiles(n=7)
+ true_vals = np.zeros(7)
+ pit = pit_from_quantiles(true_vals, qmat, levels)
+ assert pit.shape == (7,)
+
+ def test_output_in_unit_interval_for_gaussian(self):
+ rng = np.random.default_rng(7)
+ qmat, levels = self._make_gaussian_quantiles(n=20)
+ true_vals = rng.normal(size=20)
+ pit = pit_from_quantiles(true_vals, qmat, levels)
+ finite = pit[np.isfinite(pit)]
+ assert np.all((finite >= 0) & (finite <= 1))
+
+ def test_multiple_rows_independent(self):
+ # Each row has different quantile values; all true values at median
+ levels = np.array([0.1, 0.5, 0.9])
+ qmat = np.array([
+ [0.0, 1.0, 2.0],
+ [10.0, 20.0, 30.0],
+ ])
+ true_vals = np.array([1.0, 20.0]) # each at the median
+ pit = pit_from_quantiles(true_vals, qmat, levels)
+ np.testing.assert_allclose(pit, [0.5, 0.5])
+
+
+# ============================================================
+# Part 2 — calibration_curve
+# ============================================================
+
+class TestCalibrationCurve:
+
+ def test_uniform_pit_gives_empirical_near_nominal(self):
+ rng = np.random.default_rng(0)
+ pit = rng.uniform(0, 1, size=20_000)
+ curve = calibration_curve(pit)
+ for nominal, empirical in curve.items():
+ assert abs(empirical - nominal) < 0.03, (
+ f"nominal={nominal:.2f}, empirical={empirical:.3f}"
+ )
+
+ def test_overconfident_model_below_diagonal(self):
+ # All PIT near 0 or 1 (U-shaped) → CIs too narrow → empirical < nominal
+ rng = np.random.default_rng(1)
+ pit = np.concatenate([
+ rng.uniform(0.0, 0.05, 500),
+ rng.uniform(0.95, 1.0, 500),
+ ])
+ curve = calibration_curve(pit, levels=[0.9])
+ # For 90% CI we expect ~90% coverage but get ~0% (all PIT outside [0.05, 0.95])
+ assert curve[0.9] < 0.5
+
+ def test_underconfident_model_above_diagonal(self):
+ # All PIT near 0.5 → posterior spread too wide → empirical > nominal for tight CIs
+ pit = np.full(1000, 0.5)
+ curve = calibration_curve(pit, levels=[0.5])
+ # 50% CI spans [0.25, 0.75]; all PIT = 0.5 are inside → empirical = 1.0 > 0.5
+ assert curve[0.5] > 0.9
+
+ def test_returns_dict_with_default_keys(self):
+ pit = np.linspace(0.01, 0.99, 100)
+ curve = calibration_curve(pit)
+ assert set(curve.keys()) == {0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99}
+
+ def test_custom_levels(self):
+ pit = np.linspace(0, 1, 200)
+ curve = calibration_curve(pit, levels=[0.8, 0.9])
+ assert set(curve.keys()) == {0.8, 0.9}
+
+ def test_nan_in_pit_ignored(self):
+ rng = np.random.default_rng(2)
+ pit = rng.uniform(0, 1, size=1000)
+ pit_with_nan = np.concatenate([pit, np.full(50, np.nan)])
+ curve_clean = calibration_curve(pit)
+ curve_nan = calibration_curve(pit_with_nan)
+ for k in curve_clean:
+ np.testing.assert_allclose(curve_clean[k], curve_nan[k])
+
+ def test_values_in_unit_interval(self):
+ pit = np.linspace(0.01, 0.99, 200)
+ curve = calibration_curve(pit)
+ for v in curve.values():
+ assert 0.0 <= v <= 1.0
+
+ def test_empty_pit_after_nan_removal(self):
+ pit = np.full(10, np.nan)
+ with pytest.warns(RuntimeWarning):
+ curve = calibration_curve(pit)
+ # np.mean of an empty slice returns NaN
+ for v in curve.values():
+ assert np.isnan(v)
+
+
+# ============================================================
+# Part 2 — pit_uniformity_test
+# ============================================================
+
+class TestPitUniformityTest:
+
+ def test_uniform_input_high_pvalue(self):
+ rng = np.random.default_rng(42)
+ pit = rng.uniform(0, 1, size=2000)
+ result = pit_uniformity_test(pit)
+ assert result["ks_pval"] > 0.05
+
+ def test_nonuniform_input_low_pvalue(self):
+ # Bimodal near 0 and 1 — clearly not uniform
+ pit = np.concatenate([np.full(200, 0.01), np.full(200, 0.99)])
+ result = pit_uniformity_test(pit)
+ assert result["ks_pval"] < 0.001
+
+ def test_returns_expected_keys(self):
+ pit = np.linspace(0, 1, 50)
+ result = pit_uniformity_test(pit)
+ assert set(result.keys()) == {"ks_stat", "ks_pval", "mean_pit", "n_vals"}
+
+ def test_empty_input_returns_nan(self):
+ result = pit_uniformity_test(np.array([]))
+ assert np.isnan(result["ks_stat"])
+ assert np.isnan(result["ks_pval"])
+ assert np.isnan(result["mean_pit"])
+ assert result["n_vals"] == 0
+
+ def test_all_nan_returns_nan(self):
+ result = pit_uniformity_test(np.full(10, np.nan))
+ assert np.isnan(result["ks_stat"])
+ assert result["n_vals"] == 0
+
+ def test_mean_pit_near_half_for_uniform(self):
+ rng = np.random.default_rng(5)
+ pit = rng.uniform(0, 1, size=5000)
+ result = pit_uniformity_test(pit)
+ assert abs(result["mean_pit"] - 0.5) < 0.05
+
+ def test_n_vals_excludes_nan(self):
+ pit = np.array([0.1, 0.5, 0.9, np.nan, np.nan])
+ result = pit_uniformity_test(pit)
+ assert result["n_vals"] == 3
+
+
+# ============================================================
+# Part 3 — plot_pit_histogram
+# ============================================================
+
+class TestPlotPitHistogram:
+
+ def test_returns_axes(self):
+ from matplotlib.axes import Axes
+ pit = np.linspace(0.05, 0.95, 50)
+ ax = plot_pit_histogram(pit)
+ assert isinstance(ax, Axes)
+
+ def test_accepts_existing_axes(self):
+ from matplotlib import pyplot as plt
+ from matplotlib.axes import Axes
+ fig, ax = plt.subplots()
+ pit = np.linspace(0.1, 0.9, 30)
+ returned = plot_pit_histogram(pit, ax=ax)
+ assert returned is ax
+ plt.close(fig)
+
+ def test_title_set(self):
+ from matplotlib import pyplot as plt
+ pit = np.linspace(0.1, 0.9, 30)
+ ax = plot_pit_histogram(pit, title="my title")
+ assert ax.get_title() == "my title"
+ plt.close(ax.get_figure())
+
+ def test_nan_filtered_no_crash(self):
+ pit = np.array([0.1, np.nan, 0.5, np.nan, 0.9])
+ ax = plot_pit_histogram(pit)
+ from matplotlib import pyplot as plt
+ plt.close(ax.get_figure())
+
+ def test_custom_bins(self):
+ from matplotlib import pyplot as plt
+ pit = np.linspace(0.01, 0.99, 100)
+ ax = plot_pit_histogram(pit, n_bins=5)
+ # Just checks it runs without error
+ plt.close(ax.get_figure())
+
+
+# ============================================================
+# Part 3 — plot_calibration_curve
+# ============================================================
+
+class TestPlotCalibrationCurve:
+
+ def test_returns_axes(self):
+ from matplotlib.axes import Axes
+ curve = {0.5: 0.48, 0.9: 0.88, 0.95: 0.93}
+ ax = plot_calibration_curve(curve)
+ assert isinstance(ax, Axes)
+
+ def test_accepts_existing_axes(self):
+ from matplotlib import pyplot as plt
+ from matplotlib.axes import Axes
+ fig, ax = plt.subplots()
+ curve = {0.5: 0.5, 0.9: 0.9}
+ returned = plot_calibration_curve(curve, ax=ax)
+ assert returned is ax
+ plt.close(fig)
+
+ def test_empty_dict_no_crash(self):
+ from matplotlib import pyplot as plt
+ ax = plot_calibration_curve({})
+ plt.close(ax.get_figure())
+
+ def test_label_appears_in_legend(self):
+ from matplotlib import pyplot as plt
+ curve = {0.5: 0.5, 0.9: 0.9}
+ ax = plot_calibration_curve(curve, label="TestModel")
+ legend_texts = [t.get_text() for t in ax.get_legend().get_texts()]
+ assert "TestModel" in legend_texts
+ plt.close(ax.get_figure())
+
+
+# ============================================================
+# Part 4 — calibration_summary
+# ============================================================
+
+def _make_calibration_inputs(n=40, seed=0):
+ """Return (true_vals, quantile_matrix, quantile_levels) for N(0,1).
+
+ Stores 9 quantile levels; true values are drawn from N(0,1) so the PIT
+ values should be approximately uniform.
+ """
+ from scipy.stats import norm
+ rng = np.random.default_rng(seed)
+ levels = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
+ true_vals = rng.normal(size=n)
+ # All observations share the same N(0,1) posterior
+ qmat = np.tile(norm.ppf(levels), (n, 1))
+ return true_vals, qmat, levels
+
+
+class TestCalibrationSummary:
+
+ def test_writes_pit_csv(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ assert os.path.exists(f"{prefix}_pit.csv")
+
+ def test_pit_csv_has_correct_columns(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ df = pd.read_csv(f"{prefix}_pit.csv")
+ assert "true_val" in df.columns
+ assert "pit" in df.columns
+
+ def test_pit_csv_has_correct_length(self, tmp_path):
+ n = 25
+ true_vals, qmat, levels = _make_calibration_inputs(n=n)
+ prefix = str(tmp_path / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ df = pd.read_csv(f"{prefix}_pit.csv")
+ assert len(df) == n
+
+ def test_writes_calibration_curve_csv(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ assert os.path.exists(f"{prefix}_calibration_curve.csv")
+
+ def test_calibration_curve_csv_columns(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ df = pd.read_csv(f"{prefix}_calibration_curve.csv")
+ assert "nominal" in df.columns
+ assert "empirical" in df.columns
+
+ def test_writes_pdf(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ assert os.path.exists(f"{prefix}_calibration.pdf")
+
+ def test_returns_dict_with_expected_keys(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ result = calibration_summary(true_vals, qmat, levels, prefix)
+ for key in ("n_vals", "ks_stat", "ks_pval", "mean_pit", "calibration_curve"):
+ assert key in result
+
+ def test_calibration_curve_in_result(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs()
+ prefix = str(tmp_path / "cal")
+ result = calibration_summary(true_vals, qmat, levels, prefix)
+ assert isinstance(result["calibration_curve"], dict)
+ assert len(result["calibration_curve"]) > 0
+
+ def test_n_vals_matches_finite_pit(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs(n=30)
+ prefix = str(tmp_path / "cal")
+ result = calibration_summary(true_vals, qmat, levels, prefix)
+ assert result["n_vals"] == 30
+
+ def test_nan_true_vals_handled_gracefully(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs(n=20)
+ true_vals[::4] = np.nan # introduce NaNs
+ prefix = str(tmp_path / "cal")
+ result = calibration_summary(true_vals, qmat, levels, prefix)
+ assert result["n_vals"] < 20
+
+ def test_label_used_in_no_crash(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs(n=10)
+ prefix = str(tmp_path / "cal")
+ # Should not raise
+ calibration_summary(true_vals, qmat, levels, prefix, label="theta_high")
+
+ def test_creates_output_directory(self, tmp_path):
+ true_vals, qmat, levels = _make_calibration_inputs(n=10)
+ prefix = str(tmp_path / "nested" / "deep" / "cal")
+ calibration_summary(true_vals, qmat, levels, prefix)
+ assert os.path.exists(str(tmp_path / "nested" / "deep"))
+
+ def test_well_calibrated_model_passes_ks(self, tmp_path):
+ # N(0,1) true values with N(0,1) posterior → PIT ~ Uniform
+ true_vals, qmat, levels = _make_calibration_inputs(n=200, seed=99)
+ prefix = str(tmp_path / "cal")
+ result = calibration_summary(true_vals, qmat, levels, prefix)
+ # With 200 samples the KS test should not flag a calibrated model
+ assert result["ks_pval"] > 0.01
+
+
+# ============================================================
+# Part 5 — _load_h5_params (migrated from test_sbc.py)
+# ============================================================
+
+class TestLoadH5Params:
+
+ def test_round_trips_arrays(self, tmp_path):
+ data = {"alpha": np.array([1.0, 2.0, 3.0]), "beta": np.array([[4.0, 5.0]])}
+ p = str(tmp_path / "test.h5")
+ _write_h5(p, data)
+ result = _load_h5_params(p)
+ np.testing.assert_allclose(result["alpha"], data["alpha"])
+ np.testing.assert_allclose(result["beta"], data["beta"])
+
+ def test_returns_all_keys(self, tmp_path):
+ data = {"x": np.ones(3), "y": np.zeros(5)}
+ p = str(tmp_path / "test.h5")
+ _write_h5(p, data)
+ result = _load_h5_params(p)
+ assert set(result.keys()) == {"x", "y"}
+
+ def test_empty_file(self, tmp_path):
+ p = str(tmp_path / "empty.h5")
+ _write_h5(p, {})
+ result = _load_h5_params(p)
+ assert result == {}
+
+
+# ============================================================
+# Part 5 — compute_sbc_ranks (migrated from test_sbc.py)
+# ============================================================
+
+class TestComputeSBCRanks:
+
+ def _make_pair(self, tmp_path, gt_val, posterior_samples, name="param"):
+ gt_path = str(tmp_path / "run_ground_truth.h5")
+ post_path = str(tmp_path / "run_posterior.h5")
+ _write_h5(gt_path, {name: np.array(gt_val)[np.newaxis, :]})
+ _write_h5(post_path, {name: np.array(posterior_samples)})
+ return gt_path, post_path
+
+ def test_rank_zero_when_gt_below_all_posterior(self, tmp_path):
+ gt_path, post_path = self._make_pair(
+ tmp_path, gt_val=[-100.0], posterior_samples=np.ones((20, 1))
+ )
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ np.testing.assert_allclose(ranks["param"], [0.0])
+
+ def test_rank_one_when_gt_above_all_posterior(self, tmp_path):
+ gt_path, post_path = self._make_pair(
+ tmp_path, gt_val=[100.0], posterior_samples=-np.ones((20, 1))
+ )
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ np.testing.assert_allclose(ranks["param"], [1.0])
+
+ def test_rank_half_when_gt_at_median(self, tmp_path):
+ post = np.arange(-10, 11, dtype=float).reshape(21, 1)
+ gt_path, post_path = self._make_pair(tmp_path, gt_val=[0.0], posterior_samples=post)
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ np.testing.assert_allclose(ranks["param"], [10 / 21])
+
+ def test_multidimensional_param(self, tmp_path):
+ rng = np.random.default_rng(0)
+ post = rng.normal(size=(100, 3))
+ gt = np.zeros((1, 3))
+ gt_path = str(tmp_path / "run_ground_truth.h5")
+ post_path = str(tmp_path / "run_posterior.h5")
+ _write_h5(gt_path, {"param": gt})
+ _write_h5(post_path, {"param": post})
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ assert ranks["param"].shape == (3,)
+ assert np.all((ranks["param"] >= 0) & (ranks["param"] <= 1))
+
+ def test_missing_posterior_param_skipped(self, tmp_path):
+ gt_path = str(tmp_path / "run_ground_truth.h5")
+ post_path = str(tmp_path / "run_posterior.h5")
+ _write_h5(gt_path, {"param_a": np.ones((1, 2)), "param_b": np.ones((1, 2))})
+ _write_h5(post_path, {"param_a": np.ones((10, 2))})
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ assert "param_a" in ranks
+ assert "param_b" not in ranks
+
+ def test_returns_empty_when_no_common_params(self, tmp_path):
+ gt_path = str(tmp_path / "run_ground_truth.h5")
+ post_path = str(tmp_path / "run_posterior.h5")
+ _write_h5(gt_path, {"a": np.ones((1, 2))})
+ _write_h5(post_path, {"b": np.ones((10, 2))})
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ assert ranks == {}
+
+ def test_rank_values_in_unit_interval(self, tmp_path):
+ rng = np.random.default_rng(42)
+ post = rng.normal(size=(50, 4))
+ gt = rng.normal(size=(1, 4))
+ gt_path = str(tmp_path / "run_ground_truth.h5")
+ post_path = str(tmp_path / "run_posterior.h5")
+ _write_h5(gt_path, {"theta": gt})
+ _write_h5(post_path, {"theta": post})
+ ranks = compute_sbc_ranks(gt_path, post_path)
+ assert np.all((ranks["theta"] >= 0) & (ranks["theta"] <= 1))
+
+
+# ============================================================
+# Part 5 — _find_pairs (migrated from test_sbc.py)
+# ============================================================
+
+class TestFindPairs:
+
+ def test_finds_paired_files(self, tmp_path):
+ for prefix in ("run001", "run002"):
+ (tmp_path / f"{prefix}_ground_truth.h5").touch()
+ (tmp_path / f"{prefix}_posterior.h5").touch()
+ pairs = _find_pairs(str(tmp_path))
+ assert len(pairs) == 2
+ run_ids = {p[0] for p in pairs}
+ assert run_ids == {"run001", "run002"}
+
+ def test_missing_posterior_gives_none(self, tmp_path):
+ (tmp_path / "run001_ground_truth.h5").touch()
+ pairs = _find_pairs(str(tmp_path))
+ assert len(pairs) == 1
+ run_id, gt_path, post_path = pairs[0]
+ assert run_id == "run001"
+ assert post_path is None
+
+ def test_returns_empty_when_no_gt_files(self, tmp_path):
+ (tmp_path / "run001_posterior.h5").touch()
+ pairs = _find_pairs(str(tmp_path))
+ assert pairs == []
+
+ def test_gt_path_is_absolute(self, tmp_path):
+ (tmp_path / "run001_ground_truth.h5").touch()
+ (tmp_path / "run001_posterior.h5").touch()
+ pairs = _find_pairs(str(tmp_path))
+ _, gt_path, post_path = pairs[0]
+ assert os.path.isabs(gt_path)
+ assert os.path.isabs(post_path)
+
+
+# ============================================================
+# Part 5 — summarize_sbc (migrated from test_sbc.py)
+# ============================================================
+
+def _make_calibrated_sbc_dir(tmp_path, n_runs=20, n_samples=200, seed=0):
+ """Create an SBC directory with approximately uniform ranks.
+
+ For each run: gt ~ N(0,1), posterior ~ N(0,1) — ranks are Uniform(0,1).
+ """
+ rng = np.random.default_rng(seed)
+ sbc_dir = tmp_path / "sbc"
+ sbc_dir.mkdir()
+ for i in range(n_runs):
+ gt = rng.normal(size=(1, 3))
+ post = rng.normal(size=(n_samples, 3))
+ run_id = f"run{i:04d}"
+ _write_h5(str(sbc_dir / f"{run_id}_ground_truth.h5"), {"alpha": gt})
+ _write_h5(str(sbc_dir / f"{run_id}_posterior.h5"), {"alpha": post})
+ return str(sbc_dir)
+
+
+class TestSummarizeSBC:
+
+ def test_returns_dataframe_with_expected_columns(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path)
+ df = summarize_sbc(sbc_dir)
+ assert not df.empty
+ for col in ("param", "n_runs", "n_ranks", "mean_rank", "ks_stat", "ks_pval"):
+ assert col in df.columns
+
+ def test_one_row_per_parameter(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path)
+ df = summarize_sbc(sbc_dir)
+ assert len(df) == 1
+ assert df.iloc[0]["param"] == "alpha"
+
+ def test_n_runs_correct(self, tmp_path):
+ n_runs = 15
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path, n_runs=n_runs)
+ df = summarize_sbc(sbc_dir)
+ assert df.iloc[0]["n_runs"] == n_runs
+
+ def test_n_ranks_correct(self, tmp_path):
+ # 10 runs × 3 elements each = 30 rank values
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path, n_runs=10)
+ df = summarize_sbc(sbc_dir)
+ assert df.iloc[0]["n_ranks"] == 30
+
+ def test_writes_summary_csv(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path)
+ summarize_sbc(sbc_dir)
+ assert os.path.exists(os.path.join(sbc_dir, "sbc_sbc_summary.csv"))
+
+ def test_writes_ranks_csv(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path)
+ summarize_sbc(sbc_dir)
+ assert os.path.exists(os.path.join(sbc_dir, "sbc_sbc_ranks.csv"))
+
+ def test_writes_histogram_pdf(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path)
+ summarize_sbc(sbc_dir)
+ assert os.path.exists(os.path.join(sbc_dir, "sbc_rank_hist.pdf"))
+
+ def test_custom_out_prefix(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path)
+ prefix = str(tmp_path / "myprefix")
+ summarize_sbc(sbc_dir, out_prefix=prefix)
+ assert os.path.exists(f"{prefix}_sbc_summary.csv")
+ assert os.path.exists(f"{prefix}_sbc_ranks.csv")
+
+ def test_empty_dir_returns_empty_dataframe(self, tmp_path):
+ sbc_dir = str(tmp_path / "empty")
+ os.makedirs(sbc_dir)
+ df = summarize_sbc(sbc_dir)
+ assert df.empty
+
+ def test_missing_posterior_skipped_without_crash(self, tmp_path):
+ sbc_dir = tmp_path / "sbc"
+ sbc_dir.mkdir()
+ rng = np.random.default_rng(1)
+ gt = rng.normal(size=(1, 2))
+ post = rng.normal(size=(50, 2))
+ _write_h5(str(sbc_dir / "run0000_ground_truth.h5"), {"alpha": gt})
+ _write_h5(str(sbc_dir / "run0000_posterior.h5"), {"alpha": post})
+ _write_h5(str(sbc_dir / "run0001_ground_truth.h5"),
+ {"alpha": rng.normal(size=(1, 2))})
+ df = summarize_sbc(str(sbc_dir))
+ assert not df.empty
+ assert df.iloc[0]["n_runs"] == 1
+
+ def test_nonexistent_dir_raises(self, tmp_path):
+ with pytest.raises(FileNotFoundError):
+ summarize_sbc(str(tmp_path / "does_not_exist"))
+
+ def test_mean_rank_near_half_for_calibrated_model(self, tmp_path):
+ sbc_dir = _make_calibrated_sbc_dir(tmp_path, n_runs=100, n_samples=500)
+ df = summarize_sbc(sbc_dir)
+ assert abs(df.iloc[0]["mean_rank"] - 0.5) < 0.05
diff --git a/tests/tfscreen/tfmodel/components/growth_noise/__init__.py b/tests/tfscreen/tfmodel/generative/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_noise/__init__.py
rename to tests/tfscreen/tfmodel/generative/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/__init__.py b/tests/tfscreen/tfmodel/generative/components/activity/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/activity/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/activity/test_fixed.py b/tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/activity/test_fixed.py
rename to tests/tfscreen/tfmodel/generative/components/activity/test_fixed.py
diff --git a/tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py b/tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/activity/test_hierarchical_geno.py
rename to tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_geno.py
diff --git a/tests/tfscreen/tfmodel/components/activity/test_hierarchical_mut.py b/tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_mut.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/activity/test_hierarchical_mut.py
rename to tests/tfscreen/tfmodel/generative/components/activity/test_hierarchical_mut.py
diff --git a/tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py b/tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/activity/test_horseshoe_geno.py
rename to tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_geno.py
diff --git a/tests/tfscreen/tfmodel/components/activity/test_horseshoe_mut.py b/tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_mut.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/activity/test_horseshoe_mut.py
rename to tests/tfscreen/tfmodel/generative/components/activity/test_horseshoe_mut.py
diff --git a/tests/tfscreen/tfmodel/components/ln_cfu0/__init__.py b/tests/tfscreen/tfmodel/generative/components/dk_geno/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/ln_cfu0/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/dk_geno/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py b/tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/dk_geno/test_fixed.py
rename to tests/tfscreen/tfmodel/generative/components/dk_geno/test_fixed.py
diff --git a/tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py b/tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/dk_geno/test_hierarchical_geno.py
rename to tests/tfscreen/tfmodel/generative/components/dk_geno/test_hierarchical_geno.py
diff --git a/tests/tfscreen/tfmodel/components/noise/__init__.py b/tests/tfscreen/tfmodel/generative/components/growth/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/noise/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/growth/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/growth/test_linear.py b/tests/tfscreen/tfmodel/generative/components/growth/test_linear.py
similarity index 54%
rename from tests/tfscreen/tfmodel/components/growth/test_linear.py
rename to tests/tfscreen/tfmodel/generative/components/growth/test_linear.py
index d8caa731..34ce0cab 100644
--- a/tests/tfscreen/tfmodel/components/growth/test_linear.py
+++ b/tests/tfscreen/tfmodel/generative/components/growth/test_linear.py
@@ -11,6 +11,7 @@
get_hyperparameters,
get_guesses,
get_priors,
+ _parse_condition_label,
LinearParams,
)
@@ -113,14 +114,19 @@ def test_get_hyperparameters():
assert "k_scale" in params
assert "m_loc" in params
assert "m_scale" in params
- assert params["k_loc"] == 0.025
+ assert "m_scale_minus" in params
+ assert "m_scale_plus" in params
+ assert params["k_loc"] == 0.020
+ assert params["m_scale_minus"] < params["m_scale_plus"] # minus is tighter
-def test_get_priors():
+def test_get_priors_no_labels_legacy_behavior():
+ """With no condition_labels, m_is_selection is None (backward compat)."""
priors = get_priors()
assert isinstance(priors, ModelPriors)
- assert priors.k_loc == 0.025
+ assert priors.k_loc == 0.020
assert priors.m_loc == 0.0
+ assert priors.m_is_selection is None
assert not hasattr(priors, "pinned")
@@ -300,3 +306,195 @@ def test_model_and_guide_have_compatible_sample_sites(mock_data):
f" model only: {model_samples - guide_samples}\n"
f" guide only: {guide_samples - model_samples}"
)
+
+
+# ---------------------------------------------------------------------------
+# _parse_condition_label
+# ---------------------------------------------------------------------------
+
+class TestParseConditionLabel:
+
+ def test_plus_returns_true(self):
+ assert _parse_condition_label("kanR+kan") is True
+
+ def test_minus_returns_false(self):
+ assert _parse_condition_label("kanR-kan") is False
+
+ def test_selection_with_description_plus(self):
+ assert _parse_condition_label("pheS+4CP") is True
+
+ def test_control_with_description_minus(self):
+ assert _parse_condition_label("pheS-4CP") is False
+
+ def test_no_plus_no_minus_raises(self):
+ with pytest.raises(ValueError, match="does not unambiguously"):
+ _parse_condition_label("kanR_no_selection")
+
+ def test_both_plus_and_minus_raises(self):
+ with pytest.raises(ValueError, match="does not unambiguously"):
+ _parse_condition_label("weird+condition-name")
+
+ def test_error_includes_label_name(self):
+ label = "mystery_condition"
+ with pytest.raises(ValueError, match=label):
+ _parse_condition_label(label)
+
+ def test_error_explains_convention(self):
+ with pytest.raises(ValueError, match=r"\+.*selection"):
+ _parse_condition_label("unlabeled")
+
+
+# ---------------------------------------------------------------------------
+# get_priors with condition_labels
+# ---------------------------------------------------------------------------
+
+class TestGetPriorsWithLabels:
+
+ def test_plus_conditions_marked_selection(self):
+ labels = ["kanR+kan", "pheS+4CP"]
+ priors = get_priors(condition_labels=labels)
+ assert priors.m_is_selection == (True, True)
+
+ def test_minus_conditions_marked_control(self):
+ labels = ["kanR-kan", "pheS-4CP"]
+ priors = get_priors(condition_labels=labels)
+ assert priors.m_is_selection == (False, False)
+
+ def test_mixed_labels_correct_classification(self):
+ labels = ["kanR+kan", "kanR-kan", "pheS+4CP", "pheS-4CP"]
+ priors = get_priors(condition_labels=labels)
+ assert priors.m_is_selection == (True, False, True, False)
+
+ def test_single_selection_label(self):
+ priors = get_priors(condition_labels=["sel+media"])
+ assert priors.m_is_selection == (True,)
+
+ def test_invalid_label_propagates_error(self):
+ with pytest.raises(ValueError, match="unlabeled"):
+ get_priors(condition_labels=["kanR+kan", "unlabeled"])
+
+ def test_none_labels_gives_none_is_selection(self):
+ priors = get_priors(condition_labels=None)
+ assert priors.m_is_selection is None
+
+ def test_empty_labels_gives_empty_tuple(self):
+ priors = get_priors(condition_labels=[])
+ assert priors.m_is_selection == ()
+
+ def test_scales_preserved_in_priors(self):
+ priors = get_priors(condition_labels=["a+b", "a-b"])
+ assert priors.m_scale_minus == pytest.approx(0.001)
+ assert priors.m_scale_plus == pytest.approx(0.01)
+
+
+# ---------------------------------------------------------------------------
+# define_model — per-condition m prior
+# ---------------------------------------------------------------------------
+
+class TestDefineModelSelectionAware:
+
+ def _make_data(self, num_condition_rep):
+ """Minimal mock with num_condition_rep conditions."""
+ num_rep = 1
+ num_cond_pre = 1
+ num_cond_sel = num_condition_rep
+ num_tname = 1
+ num_tconc = 1
+ num_geno = 1
+ num_time = 2
+ shape = (num_rep, num_time, num_cond_pre, num_cond_sel,
+ num_tname, num_tconc, num_geno)
+
+ return MockGrowthData(
+ num_condition_rep=num_condition_rep,
+ num_replicate=num_rep,
+ map_condition_pre=jnp.zeros(num_condition_rep, dtype=jnp.int32),
+ map_condition_sel=jnp.arange(num_condition_rep, dtype=jnp.int32),
+ ln_cfu=jnp.zeros(shape),
+ t_sel=jnp.zeros(shape),
+ good_mask=jnp.ones(shape, dtype=bool),
+ )
+
+ def test_selection_aware_prior_tighter_for_control(self):
+ """
+ The Normal prior on m for a '-' condition must be tighter than for
+ a '+' condition. We verify this by inspecting the log-prob of a
+ moderate m value under each condition's prior.
+ """
+ labels = ["kanR+kan", "kanR-kan"]
+ priors = get_priors(condition_labels=labels)
+ data = self._make_data(num_condition_rep=2)
+
+ # Substitute specific m values and inspect the model trace
+ m_test_value = 0.005 # moderate m, within normal range but 5× tight range
+ subs = {
+ "test_m_k": jnp.zeros(2),
+ "test_m_m": jnp.full(2, m_test_value),
+ }
+ substituted = substitute(define_model, data=subs)
+
+ with seed(rng_seed=0):
+ tr = trace(substituted).get_trace(
+ name="test_m", data=data, priors=priors
+ )
+
+ m_site = tr["test_m_m"]
+ # The trace fn should have two m values; we can check the log-prob
+ # is lower (more penalized) for the control condition.
+ m_values = m_site["value"]
+ assert m_values.shape == (2,)
+ # kanR+kan (index 0) is selection: wider prior → less penalty
+ # kanR-kan (index 1) is control: tight prior → more penalty
+ import numpyro.distributions as d
+ lp_sel = d.Normal(priors.m_loc, priors.m_scale_plus).log_prob(m_test_value)
+ lp_ctrl = d.Normal(priors.m_loc, priors.m_scale_minus).log_prob(m_test_value)
+ assert float(lp_sel) > float(lp_ctrl)
+
+ def test_legacy_path_unchanged(self):
+ """With m_is_selection=None, the single m_scale prior applies."""
+ priors = get_priors() # no condition_labels → m_is_selection=None
+ assert priors.m_is_selection is None
+
+ data = self._make_data(num_condition_rep=2)
+ with seed(rng_seed=0):
+ tr = trace(define_model).get_trace(
+ name="leg", data=data, priors=priors
+ )
+ assert "leg_m" in tr
+ assert tr["leg_m"]["value"].shape == (2,)
+
+ def test_define_model_shapes_with_labels(self):
+ """Output shapes should be the same as the legacy path."""
+ labels = ["sel+a", "ctrl-b", "sel+c"]
+ priors = get_priors(condition_labels=labels)
+ data = self._make_data(num_condition_rep=3)
+
+ with seed(rng_seed=0):
+ params = define_model(name="sa", data=data, priors=priors)
+
+ assert params.k_pre.shape == data.map_condition_pre.shape
+ assert params.m_pre.shape == data.map_condition_pre.shape
+ assert params.k_sel.shape == data.map_condition_sel.shape
+ assert params.m_sel.shape == data.map_condition_sel.shape
+
+ def test_model_guide_compatible_with_labels(self):
+ """Model and guide sample sites must still match with labels."""
+ labels = ["kanR+kan", "kanR-kan"]
+ priors = get_priors(condition_labels=labels)
+ data = self._make_data(num_condition_rep=2)
+
+ with seed(rng_seed=0):
+ model_trace = trace(define_model).get_trace(
+ name="compat2", data=data, priors=priors
+ )
+ with seed(rng_seed=0):
+ guide_trace = trace(guide).get_trace(
+ name="compat2", data=data, priors=priors
+ )
+
+ model_samples = {
+ n for n, s in model_trace.items()
+ if s["type"] == "sample" and not s.get("is_observed", False)
+ }
+ guide_samples = {n for n, s in guide_trace.items() if s["type"] == "sample"}
+ assert model_samples == guide_samples
diff --git a/tests/tfscreen/tfmodel/components/growth/test_power.py b/tests/tfscreen/tfmodel/generative/components/growth/test_power.py
similarity index 98%
rename from tests/tfscreen/tfmodel/components/growth/test_power.py
rename to tests/tfscreen/tfmodel/generative/components/growth/test_power.py
index 937eceba..1614d9a4 100644
--- a/tests/tfscreen/tfmodel/components/growth/test_power.py
+++ b/tests/tfscreen/tfmodel/generative/components/growth/test_power.py
@@ -39,14 +39,14 @@ def test_get_hyperparameters():
assert isinstance(params, dict)
assert "k_loc" in params
assert "n_loc" in params
- assert params["k_loc"] == 0.025
+ assert params["k_loc"] == 0.020
assert params["n_loc"] == 0.0
def test_get_priors():
priors = get_priors()
assert isinstance(priors, ModelPriors)
- assert priors.k_loc == 0.025
+ assert priors.k_loc == 0.020
assert priors.n_loc == 0.0
assert not hasattr(priors, "pinned")
diff --git a/tests/tfscreen/tfmodel/components/growth/test_saturation.py b/tests/tfscreen/tfmodel/generative/components/growth/test_saturation.py
similarity index 98%
rename from tests/tfscreen/tfmodel/components/growth/test_saturation.py
rename to tests/tfscreen/tfmodel/generative/components/growth/test_saturation.py
index 3d1a7c20..2d6b7a2b 100644
--- a/tests/tfscreen/tfmodel/components/growth/test_saturation.py
+++ b/tests/tfscreen/tfmodel/generative/components/growth/test_saturation.py
@@ -39,13 +39,13 @@ def test_get_hyperparameters():
assert isinstance(params, dict)
assert "min_loc" in params
assert "max_loc" in params
- assert params["min_loc"] == 0.025
+ assert params["min_loc"] == 0.020
def test_get_priors():
priors = get_priors()
assert isinstance(priors, ModelPriors)
- assert priors.min_loc == 0.025
+ assert priors.min_loc == 0.020
assert not hasattr(priors, "pinned")
diff --git a/tests/tfscreen/tfmodel/components/sample_offset/__init__.py b/tests/tfscreen/tfmodel/generative/components/growth_noise/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/sample_offset/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/growth_noise/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/growth_noise/test_normal_kt.py b/tests/tfscreen/tfmodel/generative/components/growth_noise/test_normal_kt.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_noise/test_normal_kt.py
rename to tests/tfscreen/tfmodel/generative/components/growth_noise/test_normal_kt.py
diff --git a/tests/tfscreen/tfmodel/components/growth_noise/test_zero.py b/tests/tfscreen/tfmodel/generative/components/growth_noise/test_zero.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_noise/test_zero.py
rename to tests/tfscreen/tfmodel/generative/components/growth_noise/test_zero.py
diff --git a/tests/tfscreen/tfmodel/components/theta/__init__.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/test_baranyi.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/test_baranyi.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/test_baranyi.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/test_baranyi.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/test_baranyi_k.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/test_baranyi_k.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/test_baranyi_k.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/test_baranyi_k.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/test_baranyi_tau.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/test_baranyi_tau.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/test_baranyi_tau.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/test_baranyi_tau.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/test_instant.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/test_instant.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/test_instant.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/test_memory.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/test_memory.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/test_memory.py
diff --git a/tests/tfscreen/tfmodel/components/growth_transition/test_two_pop.py b/tests/tfscreen/tfmodel/generative/components/growth_transition/test_two_pop.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/growth_transition/test_two_pop.py
rename to tests/tfscreen/tfmodel/generative/components/growth_transition/test_two_pop.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/__init__.py b/tests/tfscreen/tfmodel/generative/components/ln_cfu0/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/ln_cfu0/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py b/tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py
similarity index 99%
rename from tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py
rename to tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py
index 1fd7db70..88970f33 100644
--- a/tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical.py
+++ b/tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical.py
@@ -325,7 +325,7 @@ def test_get_priors_no_data_returns_defaults():
def test_get_priors_accepts_data_keyword():
- """get_priors must accept an optional `data` keyword (model_class detects this)."""
+ """get_priors must accept an optional `data` keyword (ModelOrchestrator detects this)."""
sig = inspect.signature(get_priors)
assert "data" in sig.parameters
assert sig.parameters["data"].default is None
diff --git a/tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py b/tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/ln_cfu0/test_hierarchical_factored.py
rename to tests/tfscreen/tfmodel/generative/components/ln_cfu0/test_hierarchical_factored.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/__init__.py b/tests/tfscreen/tfmodel/generative/components/noise/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/noise/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/noise/test_beta.py b/tests/tfscreen/tfmodel/generative/components/noise/test_beta.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/noise/test_beta.py
rename to tests/tfscreen/tfmodel/generative/components/noise/test_beta.py
diff --git a/tests/tfscreen/tfmodel/components/noise/test_logit_normal.py b/tests/tfscreen/tfmodel/generative/components/noise/test_logit_normal.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/noise/test_logit_normal.py
rename to tests/tfscreen/tfmodel/generative/components/noise/test_logit_normal.py
diff --git a/tests/tfscreen/tfmodel/components/noise/test_zero.py b/tests/tfscreen/tfmodel/generative/components/noise/test_zero.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/noise/test_zero.py
rename to tests/tfscreen/tfmodel/generative/components/noise/test_zero.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/__init__.py b/tests/tfscreen/tfmodel/generative/components/sample_offset/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/sample_offset/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/sample_offset/test_normal.py b/tests/tfscreen/tfmodel/generative/components/sample_offset/test_normal.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/sample_offset/test_normal.py
rename to tests/tfscreen/tfmodel/generative/components/sample_offset/test_normal.py
diff --git a/tests/tfscreen/tfmodel/components/sample_offset/test_zero.py b/tests/tfscreen/tfmodel/generative/components/sample_offset/test_zero.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/sample_offset/test_zero.py
rename to tests/tfscreen/tfmodel/generative/components/sample_offset/test_zero.py
diff --git a/tests/tfscreen/tfmodel/components/test_pinning.py b/tests/tfscreen/tfmodel/generative/components/test_pinning.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/test_pinning.py
rename to tests/tfscreen/tfmodel/generative/components/test_pinning.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/theta/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/theta/test__simple.py b/tests/tfscreen/tfmodel/generative/components/theta/test__simple.py
similarity index 63%
rename from tests/tfscreen/tfmodel/components/theta/test__simple.py
rename to tests/tfscreen/tfmodel/generative/components/theta/test__simple.py
index f5555fe2..cda478a6 100644
--- a/tests/tfscreen/tfmodel/components/theta/test__simple.py
+++ b/tests/tfscreen/tfmodel/generative/components/theta/test__simple.py
@@ -54,7 +54,7 @@ def mock_data():
def mock_priors():
"""
A non-trivial theta_values tensor of shape (num_titrant_name=2,
- num_titrant_conc=3) with two saturating Hill-like rows.
+ num_titrant_conc=3) with two saturating Hill-like rows. 2-D path.
"""
theta_values = jnp.array([
[0.10, 0.50, 0.90],
@@ -63,6 +63,29 @@ def mock_priors():
return ModelPriors(theta_values=theta_values, sigma_floor=0.05)
+@pytest.fixture
+def mock_priors_3d(mock_data):
+ """
+ Per-genotype theta_values of shape (T=2, C=3, G=4).
+
+ Each genotype has a distinct sigmoid-like curve so we can verify that
+ values are NOT broadcast across the genotype axis.
+ Axis order: (titrant_name, titrant_conc, genotype).
+ """
+ # shape = (T=2, C=3, G=4): outer=titrant, middle=conc, inner=genotype.
+ theta_values = jnp.array([
+ # titrant 0 — shape (C=3, G=4)
+ [[0.90, 0.80, 0.70, 0.60], # conc 0, genotypes 0-3
+ [0.50, 0.40, 0.30, 0.20], # conc 1
+ [0.10, 0.05, 0.08, 0.03]], # conc 2
+ # titrant 1
+ [[0.85, 0.75, 0.65, 0.55],
+ [0.45, 0.35, 0.25, 0.15],
+ [0.12, 0.07, 0.09, 0.04]],
+ ]) # shape (2, 3, 4) = (T, C, G)
+ return ModelPriors(theta_values=theta_values, sigma_floor=0.05)
+
+
# ---------------------------------------------------------------------------
# Helper / sanity tests
# ---------------------------------------------------------------------------
@@ -342,3 +365,115 @@ def test_extreme_theta_values_finite_mu(mock_data):
assert jnp.isclose(theta_param.mu[0, 1, 0], 0.0)
assert theta_param.mu[0, 0, 0] < -10.0
assert theta_param.mu[0, 2, 0] > 10.0
+
+
+# ---------------------------------------------------------------------------
+# 3-D theta_values: per-genotype path
+# ---------------------------------------------------------------------------
+
+class TestPerGenotype:
+ """Tests for the (T, C, G) theta_values path introduced for calibration."""
+
+ def test_define_model_returns_per_genotype_theta(self, mock_data, mock_priors_3d):
+ """Each genotype should get its own theta column, not a broadcast copy."""
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ assert theta_param.theta.shape == (
+ mock_data.num_titrant_name,
+ mock_data.num_titrant_conc,
+ mock_data.num_genotype,
+ )
+ # The per-genotype theta must match the input exactly.
+ assert jnp.allclose(theta_param.theta, mock_priors_3d.theta_values)
+ # Genotypes must differ from each other — no broadcast collapse.
+ assert not jnp.allclose(theta_param.theta[..., 0], theta_param.theta[..., 1])
+
+ def test_define_model_3d_population_moment_shapes(self, mock_data, mock_priors_3d):
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ expected = (mock_data.num_titrant_name, mock_data.num_titrant_conc, 1)
+ assert theta_param.mu.shape == expected
+ assert theta_param.sigma.shape == expected
+
+ def test_define_model_3d_mu_is_mean_over_genotypes(self, mock_data, mock_priors_3d):
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ expected_mu = jnp.mean(
+ _logit(mock_priors_3d.theta_values), axis=-1, keepdims=True
+ )
+ assert jnp.allclose(theta_param.mu, expected_mu, atol=1e-5)
+
+ def test_define_model_3d_sigma_at_least_floor(self, mock_data, mock_priors_3d):
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ assert jnp.all(theta_param.sigma >= mock_priors_3d.sigma_floor)
+
+ def test_define_model_3d_sigma_is_floor_when_genotypes_identical(self, mock_data):
+ """When all genotypes share the same theta, sigma collapses to the floor."""
+ tv = jnp.broadcast_to(
+ jnp.array([[0.3, 0.5, 0.7], [0.2, 0.4, 0.6]])[..., None],
+ (2, 3, mock_data.num_genotype),
+ )
+ floor = 0.05
+ priors = ModelPriors(theta_values=tv, sigma_floor=floor)
+ theta_param = define_model(name="pg", data=mock_data, priors=priors)
+ assert jnp.allclose(theta_param.sigma, floor, atol=1e-5)
+
+ def test_define_model_3d_registers_no_sample_sites(self, mock_data, mock_priors_3d):
+ with seed(rng_seed=0):
+ tr = trace(define_model).get_trace(
+ name="pg", data=mock_data, priors=mock_priors_3d,
+ )
+ sample_sites = [n for n, s in tr.items() if s["type"] == "sample"]
+ assert sample_sites == []
+ assert "pg_theta" in tr
+
+ def test_guide_3d_matches_model(self, mock_data, mock_priors_3d):
+ model_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ guide_param = guide(name="pg", data=mock_data, priors=mock_priors_3d)
+ assert jnp.allclose(model_param.theta, guide_param.theta)
+ assert jnp.allclose(model_param.mu, guide_param.mu)
+
+ def test_run_model_3d_slices_correct_genotypes(self, mock_data, mock_priors_3d):
+ """run_model must pick the right per-genotype curve using geno_theta_idx."""
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ out = run_model(theta_param, mock_data._replace(scatter_theta=0))
+
+ expected_shape = (
+ mock_data.num_titrant_name,
+ mock_data.num_titrant_conc,
+ mock_data.geno_theta_idx.shape[0],
+ )
+ assert out.shape == expected_shape
+ # Each output column must match the theta_values column for that genotype.
+ for col, geno_idx in enumerate(mock_data.geno_theta_idx.tolist()):
+ assert jnp.allclose(out[..., col],
+ mock_priors_3d.theta_values[..., geno_idx])
+
+ def test_run_model_3d_concentration_mapping(self, mock_data, mock_priors_3d):
+ """Concentration remapping must work correctly with per-genotype theta."""
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ # Request concentrations [1.0, 0.0] — subset + reorder of [0, 1, 10].
+ new_conc = jnp.array([1.0, 0.0])
+ out = run_model(
+ theta_param,
+ mock_data._replace(titrant_conc=new_conc, scatter_theta=0),
+ )
+ # Expected: columns [1, 0] from theta_values, then slice by geno_theta_idx.
+ expected = mock_priors_3d.theta_values[..., mock_data.geno_theta_idx][:, [1, 0], :]
+ assert jnp.allclose(out, expected)
+
+ def test_run_model_3d_scatter(self, mock_data, mock_priors_3d):
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ out = run_model(theta_param, mock_data)
+ expected_shape = (
+ 1, 1, 1, 1,
+ mock_data.num_titrant_name,
+ mock_data.num_titrant_conc,
+ mock_data.geno_theta_idx.shape[0],
+ )
+ assert out.shape == expected_shape
+
+ def test_get_population_moments_3d(self, mock_data, mock_priors_3d):
+ theta_param = define_model(name="pg", data=mock_data, priors=mock_priors_3d)
+ mu, sigma = get_population_moments(theta_param, mock_data)
+ expected_shape = (mock_data.num_titrant_name, mock_data.num_titrant_conc, 1)
+ assert mu.shape == expected_shape
+ assert sigma.shape == expected_shape
+ assert jnp.all(sigma > 0)
diff --git a/tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py b/tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/test_categorical_geno.py
rename to tests/tfscreen/tfmodel/generative/components/theta/test_categorical_geno.py
diff --git a/tests/tfscreen/tfmodel/components/theta/test_hill_geno.py b/tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py
similarity index 75%
rename from tests/tfscreen/tfmodel/components/theta/test_hill_geno.py
rename to tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py
index e9d127f0..0fa74c2c 100644
--- a/tests/tfscreen/tfmodel/components/theta/test_hill_geno.py
+++ b/tests/tfscreen/tfmodel/generative/components/theta/test_hill_geno.py
@@ -6,14 +6,19 @@
from collections import namedtuple
# --- Import Module Under Test (MUT) ---
+import jax
+
from tfscreen.tfmodel.generative.components.theta.hill_geno import (
ModelPriors,
ThetaParam,
+ SimPriors,
define_model,
guide,
+ simulate,
run_model,
get_population_moments,
get_hyperparameters,
+ get_sim_hyperparameters,
get_guesses,
get_priors,
_population_moments,
@@ -487,3 +492,144 @@ def test_different_titrants_get_different_predictions(self):
iptg_pred = result.loc[result["titrant_name"] == "IPTG", "median"].values[0]
tmaipp_pred = result.loc[result["titrant_name"] == "TMAIPP", "median"].values[0]
assert abs(iptg_pred - tmaipp_pred) > 0.05
+
+
+# ---------------------------------------------------------------------------
+# SimPriors / get_sim_hyperparameters
+# ---------------------------------------------------------------------------
+
+def test_get_sim_hyperparameters_returns_dict():
+ params = get_sim_hyperparameters()
+ assert isinstance(params, dict)
+ for key in ["wt_theta_low", "wt_theta_high", "wt_log_K", "wt_hill_n",
+ "sigma_logit_low", "sigma_logit_delta", "sigma_log_K", "sigma_log_n",
+ "p_stuck_bound", "p_never_binds", "p_inverted"]:
+ assert key in params
+
+
+def test_sim_priors_constructs():
+ sp = SimPriors(**get_sim_hyperparameters())
+ assert isinstance(sp, SimPriors)
+ assert 0.0 < sp.wt_theta_low < 1.0
+ assert 0.0 < sp.wt_theta_high < 1.0
+ assert sp.p_stuck_bound + sp.p_never_binds + sp.p_inverted < 1.0
+
+
+# ---------------------------------------------------------------------------
+# simulate
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def sim_mock_data(mock_data):
+ """Mock data compatible with simulate() — only needs num_genotype and log_titrant_conc."""
+ return mock_data._replace(scatter_theta=0)
+
+
+def test_simulate_output_shapes(sim_mock_data):
+ rng_key = jax.random.PRNGKey(0)
+ sim_priors = SimPriors(**get_sim_hyperparameters())
+ theta_gc, theta_param = simulate("theta", sim_mock_data, sim_priors, rng_key)
+
+ G = sim_mock_data.num_genotype
+ C = sim_mock_data.num_titrant_conc
+
+ assert theta_gc.shape == (G, C)
+ assert theta_param.theta_low.shape == (1, G)
+ assert theta_param.theta_high.shape == (1, G)
+ assert theta_param.log_hill_K.shape == (1, G)
+ assert theta_param.hill_n.shape == (1, G)
+ assert theta_param.mu is None
+ assert theta_param.sigma is None
+
+
+def test_simulate_theta_in_unit_interval(sim_mock_data):
+ rng_key = jax.random.PRNGKey(1)
+ sim_priors = SimPriors(**get_sim_hyperparameters())
+ theta_gc, _ = simulate("theta", sim_mock_data, sim_priors, rng_key)
+ assert np.all(theta_gc >= 0.0)
+ assert np.all(theta_gc <= 1.0)
+
+
+def test_simulate_normal_genotypes_wide_range(mock_data):
+ """With many genotypes and all normal, most curves should span > 0.5 dynamic range."""
+ WideData = namedtuple("WideData", mock_data._fields)
+ data = WideData(**{
+ **mock_data._asdict(),
+ "num_genotype": 200,
+ "log_titrant_conc": jnp.linspace(-8, 2, 20),
+ "num_titrant_conc": 20,
+ "scatter_theta": 0,
+ })
+
+ params = get_sim_hyperparameters()
+ params.update({"p_stuck_bound": 0.0, "p_never_binds": 0.0, "p_inverted": 0.0})
+ sim_priors = SimPriors(**params)
+
+ rng_key = jax.random.PRNGKey(42)
+ theta_gc, _ = simulate("theta", data, sim_priors, rng_key)
+
+ dynamic_range = np.ptp(theta_gc, axis=1) # max - min per genotype
+ assert np.median(dynamic_range) > 0.5, (
+ f"Median dynamic range {np.median(dynamic_range):.3f} is too small"
+ )
+
+
+def test_simulate_stuck_bound_genotypes(mock_data):
+ """Stuck-bound genotypes should have near-zero dynamic range."""
+ LargeData = namedtuple("LargeData", mock_data._fields)
+ data = LargeData(**{
+ **mock_data._asdict(),
+ "num_genotype": 500,
+ "log_titrant_conc": jnp.linspace(-8, 2, 20),
+ "num_titrant_conc": 20,
+ "scatter_theta": 0,
+ })
+
+ params = get_sim_hyperparameters()
+ params.update({"p_stuck_bound": 1.0, "p_never_binds": 0.0, "p_inverted": 0.0})
+ sim_priors = SimPriors(**params)
+
+ rng_key = jax.random.PRNGKey(7)
+ theta_gc, _ = simulate("theta", data, sim_priors, rng_key)
+
+ dynamic_range = np.ptp(theta_gc, axis=1)
+ assert np.all(dynamic_range < 0.3), (
+ f"Stuck-bound genotypes should have small dynamic range; max was {dynamic_range.max():.3f}"
+ )
+
+
+def test_simulate_mixture_probabilities_error(sim_mock_data):
+ """SimPriors with mixture probabilities summing > 1 should raise."""
+ params = get_sim_hyperparameters()
+ params.update({"p_stuck_bound": 0.5, "p_never_binds": 0.5, "p_inverted": 0.2})
+ sim_priors = SimPriors(**params)
+ with pytest.raises(ValueError, match="sum to more than 1"):
+ simulate("theta", sim_mock_data, sim_priors, jax.random.PRNGKey(0))
+
+
+def test_simulate_deterministic_with_same_key(sim_mock_data):
+ """Same rng_key should produce identical results."""
+ sim_priors = SimPriors(**get_sim_hyperparameters())
+ t1, _ = simulate("theta", sim_mock_data, sim_priors, jax.random.PRNGKey(99))
+ t2, _ = simulate("theta", sim_mock_data, sim_priors, jax.random.PRNGKey(99))
+ np.testing.assert_array_equal(t1, t2)
+
+
+def test_simulate_different_keys_differ(sim_mock_data):
+ """Different rng_keys should (almost always) produce different results."""
+ sim_priors = SimPriors(**get_sim_hyperparameters())
+ t1, _ = simulate("theta", sim_mock_data, sim_priors, jax.random.PRNGKey(0))
+ t2, _ = simulate("theta", sim_mock_data, sim_priors, jax.random.PRNGKey(1))
+ assert not np.allclose(t1, t2)
+
+
+def test_simulate_theta_param_compatible_with_run_model(mock_data):
+ """ThetaParam from simulate() should be usable by run_model without error."""
+ data = mock_data._replace(scatter_theta=0)
+ sim_priors = SimPriors(**get_sim_hyperparameters())
+ _, theta_param = simulate("theta", data, sim_priors, jax.random.PRNGKey(5))
+ result = run_model(theta_param, data)
+ G = data.num_genotype
+ C = data.num_titrant_conc
+ # simulate() always produces T=1 (no per-titrant structure); run_model reflects that
+ assert result.shape == (1, C, G)
diff --git a/tests/tfscreen/tfmodel/components/theta/test_hill_mut.py b/tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py
similarity index 55%
rename from tests/tfscreen/tfmodel/components/theta/test_hill_mut.py
rename to tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py
index 24244e17..d17fbc18 100644
--- a/tests/tfscreen/tfmodel/components/theta/test_hill_mut.py
+++ b/tests/tfscreen/tfmodel/generative/components/theta/test_hill_mut.py
@@ -1,3 +1,5 @@
+import warnings
+
import pytest
import numpy as np
import jax.numpy as jnp
@@ -5,15 +7,19 @@
from numpyro.handlers import trace, substitute, seed
from collections import namedtuple
+import jax
from functools import partial
from tfscreen.tfmodel.generative.components.theta.hill_mut import (
ModelPriors,
ThetaParam,
+ SimPriors,
define_model,
guide,
+ simulate,
run_model,
get_population_moments,
get_hyperparameters,
+ get_sim_hyperparameters,
get_guesses,
get_priors,
_assemble,
@@ -452,3 +458,323 @@ def test_get_population_moments_shape(mock_data_no_epi):
assert mu.shape == (T, C, 1)
assert sigma.shape == (T, C, 1)
assert jnp.all(sigma > 0)
+
+
+# ---------------------------------------------------------------------------
+# SimPriors / get_sim_hyperparameters
+# ---------------------------------------------------------------------------
+
+def test_get_sim_hyperparameters_returns_dict():
+ params = get_sim_hyperparameters()
+ assert isinstance(params, dict)
+ for key in ["wt_theta_low", "wt_theta_high", "wt_log_K", "wt_hill_n",
+ "sigma_d_logit_low", "sigma_d_logit_delta",
+ "sigma_d_log_K", "sigma_d_log_n",
+ "epi_tau_scale", "epi_slab_scale", "epi_slab_df"]:
+ assert key in params, f"Missing key: {key}"
+
+
+def test_get_sim_hyperparameters_horseshoe_defaults_positive():
+ """Horseshoe defaults must all be strictly positive."""
+ params = get_sim_hyperparameters()
+ assert params["epi_tau_scale"] > 0.0
+ assert params["epi_slab_scale"] > 0.0
+ assert params["epi_slab_df"] > 0.0
+
+
+def test_sim_priors_constructs():
+ sp = SimPriors(**get_sim_hyperparameters())
+ assert isinstance(sp, SimPriors)
+ assert 0.0 < sp.wt_theta_low < 1.0
+ assert 0.0 < sp.wt_theta_high < 1.0
+ assert sp.wt_hill_n > 0.0
+ # Horseshoe attributes must exist and be non-negative
+ assert sp.epi_tau_scale >= 0.0
+ assert sp.epi_slab_scale >= 0.0
+ assert sp.epi_slab_df >= 0.0
+
+
+# ---------------------------------------------------------------------------
+# simulate – no epistasis
+# ---------------------------------------------------------------------------
+
+class TestSimulateNoEpi:
+
+ def test_output_shapes(self, mock_data_no_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ theta_gc, theta_param = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(0))
+ G = mock_data_no_epi.num_genotype
+ C = mock_data_no_epi.num_titrant_conc
+ assert theta_gc.shape == (G, C)
+ assert theta_param.theta_low.shape == (1, G)
+ assert theta_param.theta_high.shape == (1, G)
+ assert theta_param.log_hill_K.shape == (1, G)
+ assert theta_param.hill_n.shape == (1, G)
+ assert theta_param.mu is None
+ assert theta_param.sigma is None
+
+ def test_theta_in_unit_interval(self, mock_data_no_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ theta_gc, _ = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(1))
+ assert np.all(theta_gc >= 0.0)
+ assert np.all(theta_gc <= 1.0)
+
+ def test_wt_genotype_exact_reference(self, mock_data_no_epi):
+ """Wildtype (column 0, no mutations) receives zero additive deltas, so
+ its assembled parameters must exactly match the WT SimPriors reference."""
+ sp = SimPriors(**get_sim_hyperparameters())
+ _, theta_param = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(2))
+ # Column 0 is WT: no mutation indicator → zero additive delta on every param
+ wt_theta_low = float(theta_param.theta_low[0, 0])
+ wt_theta_high = float(theta_param.theta_high[0, 0])
+ assert abs(wt_theta_low - sp.wt_theta_low) < 1e-5
+ assert abs(wt_theta_high - sp.wt_theta_high) < 1e-5
+ # WT curve must be decreasing: theta goes from ~1 at low conc to ~0 at high conc
+ assert wt_theta_low > wt_theta_high
+
+ def test_mutant_differs_from_wt(self, mock_data_no_epi):
+ """With non-zero sigma_d_log_K, mutant genotypes should have a shifted K."""
+ params = get_sim_hyperparameters()
+ params["sigma_d_log_K"] = 2.0 # large effect to ensure visible shift
+ sp = SimPriors(**params)
+ theta_gc, _ = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(3))
+ # Single mutants (columns 1 and 2) should differ from wt (column 0) at mid-conc
+ assert not np.allclose(theta_gc[0], theta_gc[1], atol=1e-3) or \
+ not np.allclose(theta_gc[0], theta_gc[2], atol=1e-3)
+
+ def test_deterministic_same_key(self, mock_data_no_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ t1, _ = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(99))
+ t2, _ = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(99))
+ np.testing.assert_array_equal(t1, t2)
+
+ def test_different_keys_differ(self, mock_data_no_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ t1, _ = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(0))
+ t2, _ = simulate("theta", mock_data_no_epi, sp, jax.random.PRNGKey(1))
+ assert not np.allclose(t1, t2)
+
+ def test_compatible_with_run_model(self, mock_data_no_epi):
+ """ThetaParam from simulate() should be usable by run_model."""
+ data = mock_data_no_epi._replace(scatter_theta=0)
+ sp = SimPriors(**get_sim_hyperparameters())
+ _, theta_param = simulate("theta", data, sp, jax.random.PRNGKey(5))
+ result = run_model(theta_param, data)
+ G = data.num_genotype
+ C = data.num_titrant_conc
+ # simulate() always produces T=1
+ assert result.shape == (1, C, G)
+
+
+# ---------------------------------------------------------------------------
+# simulate – with epistasis
+# ---------------------------------------------------------------------------
+
+class TestSimulateEpi:
+
+ def test_output_shapes(self, mock_data_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ theta_gc, theta_param = simulate("theta", mock_data_epi, sp, jax.random.PRNGKey(0))
+ G = mock_data_epi.num_genotype
+ C = mock_data_epi.num_titrant_conc
+ assert theta_gc.shape == (G, C)
+ assert theta_param.theta_low.shape == (1, G)
+
+ def test_theta_in_unit_interval(self, mock_data_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ theta_gc, _ = simulate("theta", mock_data_epi, sp, jax.random.PRNGKey(7))
+ assert np.all(theta_gc >= 0.0)
+ assert np.all(theta_gc <= 1.0)
+
+ def test_epistasis_only_shifts_double_mutant(self, mock_data_epi):
+ """
+ With zero mutation deltas but non-zero epistasis, only the double-mutant
+ column (index 3) can differ from the wildtype column (index 0).
+ Single mutants (indices 1 and 2) have no pair membership, so they must
+ equal the wildtype reference regardless of the epistasis draw.
+ """
+ params = get_sim_hyperparameters()
+ params.update({
+ "sigma_d_logit_low": 0.0,
+ "sigma_d_logit_delta": 0.0,
+ "sigma_d_log_K": 0.0,
+ "sigma_d_log_n": 0.0,
+ "epi_tau_scale": 2.0, # large τ → large horseshoe effects
+ "epi_slab_scale": 2.0,
+ "epi_slab_df": 4.0,
+ })
+ sp = SimPriors(**params)
+ theta_gc, _ = simulate("theta", mock_data_epi, sp, jax.random.PRNGKey(42))
+ # Columns 0, 1, 2 have no pair membership → must equal WT reference
+ np.testing.assert_allclose(theta_gc[0], theta_gc[1], atol=1e-6)
+ np.testing.assert_allclose(theta_gc[0], theta_gc[2], atol=1e-6)
+ # Column 3 (double mutant M42I/K84L) may differ due to epistasis
+
+ def test_epi_tau_zero_produces_no_epistasis(self, mock_data_epi):
+ """
+ epi_tau_scale=0.0 makes τ=0 exactly, so all epistasis effects are zero.
+ With zero mutation deltas too, all four genotypes must be identical.
+ """
+ params = get_sim_hyperparameters()
+ params.update({
+ "sigma_d_logit_low": 0.0,
+ "sigma_d_logit_delta": 0.0,
+ "sigma_d_log_K": 0.0,
+ "sigma_d_log_n": 0.0,
+ "epi_tau_scale": 0.0, # hard off-switch: τ = 0 exactly
+ })
+ sp = SimPriors(**params)
+ theta_gc, _ = simulate("theta", mock_data_epi, sp, jax.random.PRNGKey(42))
+ # No mutation delta, no epistasis → all columns are the WT reference
+ np.testing.assert_allclose(theta_gc[0], theta_gc[1], atol=1e-10)
+ np.testing.assert_allclose(theta_gc[0], theta_gc[2], atol=1e-10)
+ np.testing.assert_allclose(theta_gc[0], theta_gc[3], atol=1e-10)
+
+ def test_horseshoe_larger_tau_produces_larger_effects(self, mock_data_epi):
+ """
+ Over many random seeds, a larger epi_tau_scale should produce larger
+ median absolute epistasis effects on the double-mutant log_K.
+ """
+ base = get_sim_hyperparameters()
+ base.update({
+ "sigma_d_logit_low": 0.0,
+ "sigma_d_logit_delta": 0.0,
+ "sigma_d_log_K": 0.0,
+ "sigma_d_log_n": 0.0,
+ })
+
+ def median_abs_epi(tau_scale, n=200):
+ params = dict(base)
+ params["epi_tau_scale"] = tau_scale
+ sp = SimPriors(**params)
+ effects = []
+ for seed_val in range(n):
+ _, tp = simulate("theta", mock_data_epi, sp,
+ jax.random.PRNGKey(seed_val))
+ # Epistasis affects col 3 (double mutant); col 0 is pure WT
+ effects.append(float(tp.log_hill_K[0, 3] - tp.log_hill_K[0, 0]))
+ return float(np.median(np.abs(effects)))
+
+ with warnings.catch_warnings():
+ warnings.filterwarnings("ignore", message="overflow encountered in exp",
+ category=RuntimeWarning)
+ small_epi = median_abs_epi(0.1)
+ large_epi = median_abs_epi(1.0)
+ assert large_epi > small_epi, (
+ f"Expected larger tau to produce larger effects, but "
+ f"small={small_epi:.4f}, large={large_epi:.4f}")
+
+ def test_deterministic_same_key(self, mock_data_epi):
+ sp = SimPriors(**get_sim_hyperparameters())
+ t1, _ = simulate("theta", mock_data_epi, sp, jax.random.PRNGKey(11))
+ t2, _ = simulate("theta", mock_data_epi, sp, jax.random.PRNGKey(11))
+ np.testing.assert_array_equal(t1, t2)
+
+
+# ---------------------------------------------------------------------------
+# define_model with num_mutation == 0 (wt-only / stratified-pool case)
+# ---------------------------------------------------------------------------
+#
+# This was the bug: numpyro.plate requires size > 0, so calling define_model
+# with a wt-only library (M=0) raised AssertionError inside the mutation plate.
+# The fix guards the plate block with `if num_mut > 0` and returns zero-shape
+# offset arrays when M=0. These tests lock in that behaviour.
+
+@pytest.fixture
+def mock_data_no_mutation():
+ """1 genotype (wt only), 0 mutations, 0 pairs — the stratified pool case."""
+ return MockData(
+ num_titrant_name=1,
+ num_titrant_conc=4,
+ num_genotype=1,
+ log_titrant_conc=jnp.linspace(-5, 0, 4),
+ geno_theta_idx=jnp.array([0], dtype=jnp.int32),
+ scatter_theta=0,
+ num_mutation=0,
+ num_pair=0,
+ mut_geno_matrix=np.zeros((0, 1), dtype=np.float32),
+ mut_nnz_mut_idx=np.zeros(0, dtype=np.int32),
+ mut_nnz_geno_idx=np.zeros(0, dtype=np.int32),
+ pair_nnz_pair_idx=np.zeros(0, dtype=np.int32),
+ pair_nnz_geno_idx=np.zeros(0, dtype=np.int32),
+ batch_idx=jnp.array([0], dtype=jnp.int32),
+ )
+
+
+class TestDefineModelNoMutation:
+
+ def test_does_not_raise(self, mock_data_no_mutation):
+ """define_model must not raise when num_mutation == 0."""
+ priors = get_priors()
+ with seed(rng_seed=jax.random.PRNGKey(0)):
+ tp = define_model("theta", mock_data_no_mutation, priors)
+ assert isinstance(tp, ThetaParam)
+
+ def test_output_shapes(self, mock_data_no_mutation):
+ """ThetaParam fields must have shape (T=1, G=1) with M=0."""
+ priors = get_priors()
+ with seed(rng_seed=jax.random.PRNGKey(1)):
+ tp = define_model("theta", mock_data_no_mutation, priors)
+ T, G = mock_data_no_mutation.num_titrant_name, mock_data_no_mutation.num_genotype
+ assert tp.theta_low.shape == (T, G)
+ assert tp.theta_high.shape == (T, G)
+ assert tp.log_hill_K.shape == (T, G)
+ assert tp.hill_n.shape == (T, G)
+
+ def test_per_genotype_equals_wt(self, mock_data_no_mutation):
+ """With M=0 the single genotype column must equal the WT parameters exactly."""
+ priors = get_priors()
+ with seed(rng_seed=jax.random.PRNGKey(2)):
+ tp = define_model("theta", mock_data_no_mutation, priors)
+ # theta values must be valid probabilities
+ assert jnp.all(tp.theta_low >= 0.0) and jnp.all(tp.theta_low <= 1.0)
+ assert jnp.all(tp.theta_high >= 0.0) and jnp.all(tp.theta_high <= 1.0)
+
+ def test_different_keys_give_different_curves(self, mock_data_no_mutation):
+ """Independent rng_keys must produce different WT parameter draws."""
+ priors = get_priors()
+ with seed(rng_seed=jax.random.PRNGKey(10)):
+ tp1 = define_model("theta", mock_data_no_mutation, priors)
+ with seed(rng_seed=jax.random.PRNGKey(11)):
+ tp2 = define_model("theta", mock_data_no_mutation, priors)
+ assert not jnp.allclose(tp1.log_hill_K, tp2.log_hill_K), (
+ "Different seeds should produce different log_K draws")
+
+ def test_force_prior_predictive_path_does_not_raise(self, mock_data_no_mutation):
+ """sample_theta_prior with force_prior_predictive=True must work for M=0.
+
+ This is the exact call path used by sample_theta_stratified when building
+ the diversity pool for mutation-decomposed components like hill_mut.
+ """
+ from tfscreen.simulate.sample_theta import sample_theta_prior
+ theta_gc, theta_param = sample_theta_prior(
+ "hill_mut", mock_data_no_mutation,
+ rng_key=jax.random.PRNGKey(42),
+ force_prior_predictive=True,
+ )
+ assert theta_gc.shape == (mock_data_no_mutation.num_genotype,
+ mock_data_no_mutation.num_titrant_conc)
+
+ def test_pool_members_diverse_under_force_prior_predictive(self, mock_data_no_mutation):
+ """Multiple pool members must produce different curves (not all wt-identical).
+
+ This is the regression test for the original bug: the perturbation path
+ with M=0 returned the same fixed wt curve for every pool member.
+ """
+ from tfscreen.simulate.sample_theta import sample_theta_prior
+ import jax
+ rng_key = jax.random.PRNGKey(0)
+ curves = []
+ for i in range(10):
+ key_i = jax.random.fold_in(rng_key, i)
+ theta_gc, _ = sample_theta_prior(
+ "hill_mut", mock_data_no_mutation,
+ rng_key=key_i,
+ force_prior_predictive=True,
+ )
+ curves.append(theta_gc[0]) # single genotype → (C,)
+
+ curves = np.stack(curves) # (10, C)
+ # At least some curves must differ (pool is not homogeneous)
+ assert curves.std(axis=0).max() > 0.01, (
+ "Pool members are all identical — prior-predictive path is not sampling")
diff --git a/tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py b/tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/test_predict_unmeasured.py
rename to tests/tfscreen/tfmodel/generative/components/theta/test_predict_unmeasured.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PK.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PddG.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_PnnC.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_predict_unmeasured.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U0_a/test_thermo.py
diff --git a/tests/tfscreen/tfmodel/components/theta_rescale/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta_rescale/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PK.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PddG.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_PnnC.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C12_K5_U1_a/test_thermo.py
diff --git a/tests/tfscreen/tfmodel/components/transformation/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/__init__.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/transformation/__init__.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/__init__.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PK.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PddG.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_PnnC.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_predict_unmeasured.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U0_a/test_thermo.py
diff --git a/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PK.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PddG.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_PnnC.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/O2_C4_K3_U1_a/test_thermo.py
diff --git a/tests/tfscreen/tfmodel/generative/components/theta/thermo/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/test_features.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/test_features.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/test_features.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/test_horseshoe.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/test_horseshoe.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/test_io.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/test_io.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/test_io.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/test_nn.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/test_nn.py
diff --git a/tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py b/tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta/thermo/test_prior.py
rename to tests/tfscreen/tfmodel/generative/components/theta/thermo/test_prior.py
diff --git a/tests/tfscreen/tfmodel/generative/components/theta_rescale/__init__.py b/tests/tfscreen/tfmodel/generative/components/theta_rescale/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/tfscreen/tfmodel/components/theta_rescale/test_logit.py b/tests/tfscreen/tfmodel/generative/components/theta_rescale/test_logit.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta_rescale/test_logit.py
rename to tests/tfscreen/tfmodel/generative/components/theta_rescale/test_logit.py
diff --git a/tests/tfscreen/tfmodel/components/theta_rescale/test_passthrough.py b/tests/tfscreen/tfmodel/generative/components/theta_rescale/test_passthrough.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/theta_rescale/test_passthrough.py
rename to tests/tfscreen/tfmodel/generative/components/theta_rescale/test_passthrough.py
diff --git a/tests/tfscreen/tfmodel/generative/components/transformation/__init__.py b/tests/tfscreen/tfmodel/generative/components/transformation/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/tfscreen/tfmodel/components/transformation/test__congression.py b/tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/transformation/test__congression.py
rename to tests/tfscreen/tfmodel/generative/components/transformation/test__congression.py
diff --git a/tests/tfscreen/tfmodel/components/transformation/test_empirical.py b/tests/tfscreen/tfmodel/generative/components/transformation/test_empirical.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/transformation/test_empirical.py
rename to tests/tfscreen/tfmodel/generative/components/transformation/test_empirical.py
diff --git a/tests/tfscreen/tfmodel/components/transformation/test_logit_norm.py b/tests/tfscreen/tfmodel/generative/components/transformation/test_logit_norm.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/transformation/test_logit_norm.py
rename to tests/tfscreen/tfmodel/generative/components/transformation/test_logit_norm.py
diff --git a/tests/tfscreen/tfmodel/components/transformation/test_single.py b/tests/tfscreen/tfmodel/generative/components/transformation/test_single.py
similarity index 100%
rename from tests/tfscreen/tfmodel/components/transformation/test_single.py
rename to tests/tfscreen/tfmodel/generative/components/transformation/test_single.py
diff --git a/tests/tfscreen/tfmodel/observe/test_observe_binding.py b/tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py
similarity index 100%
rename from tests/tfscreen/tfmodel/observe/test_observe_binding.py
rename to tests/tfscreen/tfmodel/generative/observe/test_observe_binding.py
diff --git a/tests/tfscreen/tfmodel/observe/test_observe_growth.py b/tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py
similarity index 100%
rename from tests/tfscreen/tfmodel/observe/test_observe_growth.py
rename to tests/tfscreen/tfmodel/generative/observe/test_observe_growth.py
diff --git a/tests/tfscreen/tfmodel/scripts/test_extract_params_cli.py b/tests/tfscreen/tfmodel/scripts/test_extract_params_cli.py
index c79744dd..99e0cae9 100644
--- a/tests/tfscreen/tfmodel/scripts/test_extract_params_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_extract_params_cli.py
@@ -112,14 +112,15 @@ def test_creates_csv_files(self, tmp_path):
assert os.path.exists(f"{out_prefix}_activity.csv")
def test_passes_point_est_q_to_get(self, tmp_path):
+ """q_to_get=[0.5] is passed to extract_parameters for MAP checkpoints."""
model = ExtractModel(num_genotype=4)
ckpt_path = _make_map_checkpoint(tmp_path, model, step=100)
out_prefix = str(tmp_path / "out")
captured = {}
- def fake_extract(gm, posteriors, q_to_get=None):
+ def fake_extract(orchestrator, posteriors, q_to_get=None):
captured["q_to_get"] = q_to_get
- return {"activity": pd.DataFrame({"genotype": ["wt"], "point_est": [0.5]})}
+ return {"activity": pd.DataFrame({"genotype": ["wt"], "q0.5": [0.5]})}
fake_gm = FakeTFModel(model)
with patch(
@@ -133,14 +134,14 @@ def fake_extract(gm, posteriors, q_to_get=None):
):
extract_params("cfg.yaml", ckpt_path, out_prefix=out_prefix)
- assert captured["q_to_get"] == {"point_est": 0.5}
+ assert captured["q_to_get"] == [0.5]
def test_posteriors_have_leading_sample_dim(self, tmp_path):
model = ExtractModel(num_genotype=4)
ckpt_path = _make_map_checkpoint(tmp_path, model, step=100)
captured = {}
- def fake_extract(gm, posteriors, q_to_get=None):
+ def fake_extract(orchestrator, posteriors, q_to_get=None):
captured["posteriors"] = posteriors
return {"activity": pd.DataFrame({"genotype": ["wt"], "point_est": [0.5]})}
@@ -175,7 +176,7 @@ def test_calls_extract_with_default_quantiles(self, tmp_path):
with h5py.File(posterior_file, "w") as f:
f.create_dataset("dummy_param", data=np.ones((10, 4)))
- def fake_extract(gm, posteriors, q_to_get=None):
+ def fake_extract(orchestrator, posteriors, q_to_get=None):
captured["q_to_get"] = q_to_get
return {"activity": pd.DataFrame({"genotype": ["wt"], "median": [0.5]})}
diff --git a/tests/tfscreen/tfmodel/scripts/test_fit_model_cli.py b/tests/tfscreen/tfmodel/scripts/test_fit_model_cli.py
index a64c810c..fcb70140 100644
--- a/tests/tfscreen/tfmodel/scripts/test_fit_model_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_fit_model_cli.py
@@ -8,6 +8,7 @@
from tfscreen.tfmodel.scripts.fit_model_cli import (
fit_model,
+ _run_svi,
)
@@ -235,3 +236,43 @@ def test_svi_epoch_checkpoint_interval_default_is_1000(self, tmp_path, mocker):
kwargs = run_svi_mock.call_args.kwargs
assert kwargs["epoch_checkpoint_interval"] == 1000
+
+
+# ---------------------------------------------------------------------------
+# _run_svi — convergence message suppression when max_num_epochs=0
+# ---------------------------------------------------------------------------
+
+class TestRunSviConvergenceMessage:
+ """_run_svi must not print a converged/not-converged message when
+ max_num_epochs=0 (checkpoint-restore mode used by tfs-sample-posterior)."""
+
+ def _make_ri(self):
+ ri = MagicMock()
+ ri._iterations_per_epoch = 1
+ ri.setup_svi.return_value = MagicMock()
+ # run_optimization returns (state, params, converged=False)
+ ri.run_optimization.return_value = (MagicMock(), {}, False)
+ return ri
+
+ def test_no_message_when_zero_epochs(self, capsys):
+ ri = self._make_ri()
+ _run_svi(ri, init_params={}, checkpoint_file=None,
+ out_prefix="tfs", max_num_epochs=0)
+ out = capsys.readouterr().out
+ assert "converged" not in out.lower()
+
+ def test_converged_message_when_epochs_run_and_converged(self, capsys):
+ ri = self._make_ri()
+ ri.run_optimization.return_value = (MagicMock(), {}, True)
+ _run_svi(ri, init_params={}, checkpoint_file=None,
+ out_prefix="tfs", max_num_epochs=10)
+ out = capsys.readouterr().out
+ assert "SVI run converged." in out
+
+ def test_not_converged_message_when_epochs_run_and_not_converged(self, capsys):
+ ri = self._make_ri()
+ ri.run_optimization.return_value = (MagicMock(), {}, False)
+ _run_svi(ri, init_params={}, checkpoint_file=None,
+ out_prefix="tfs", max_num_epochs=10)
+ out = capsys.readouterr().out
+ assert "SVI run has not yet converged." in out
diff --git a/tests/tfscreen/tfmodel/scripts/test_predict_growth_cli.py b/tests/tfscreen/tfmodel/scripts/test_predict_growth_cli.py
index 293bd5fe..8c34a719 100644
--- a/tests/tfscreen/tfmodel/scripts/test_predict_growth_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_predict_growth_cli.py
@@ -30,34 +30,34 @@ def _write_lines(path, lines):
@pytest.fixture
-def mock_gm():
- gm = MagicMock()
- gm.growth_df = _make_growth_df(["wt", "A1B"], [0.0, 1.0])
- return gm
+def mock_orchestrator():
+ orchestrator = MagicMock()
+ orchestrator.growth_df = _make_growth_df(["wt", "A1B"], [0.0, 1.0])
+ return orchestrator
@pytest.fixture
-def mock_predict(mock_gm):
+def mock_predict(mock_orchestrator):
"""Patch read_configuration, resolve_param_file, and predict; return captured call kwargs."""
calls = {}
def fake_predict(**kwargs):
calls.update(kwargs)
- genotypes = kwargs.get("genotypes") or mock_gm.growth_df["genotype"].unique().tolist()
- concs = kwargs.get("titrant_conc") or mock_gm.growth_df["titrant_conc"].unique().tolist()
+ genotypes = kwargs.get("genotypes") or mock_orchestrator.growth_df["genotype"].unique().tolist()
+ concs = kwargs.get("titrant_conc") or mock_orchestrator.growth_df["titrant_conc"].unique().tolist()
rows = [{"genotype": g, "titrant_name": "IPTG", "titrant_conc": c,
- "median": 10.0}
+ "q0.5": 10.0}
for g in genotypes for c in concs]
return pd.DataFrame(rows)
with patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.read_configuration",
- return_value=(mock_gm, {}),
+ return_value=(mock_orchestrator, {}),
), patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.resolve_param_file",
- side_effect=lambda pf, gm, op: pf, # pass-through
+ side_effect=lambda pf, orchestrator, op: pf, # pass-through
), patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.predict",
@@ -87,7 +87,7 @@ def test_no_files_passes_none_concs(self, mock_predict, tmp_path):
class TestPredictGrowthUnion:
- def test_genotypes_file_unions_with_training(self, mock_predict, mock_gm, tmp_path):
+ def test_genotypes_file_unions_with_training(self, mock_predict, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["C2D"]) # novel genotype
predict_growth("cfg.yaml", "post.h5",
@@ -97,7 +97,7 @@ def test_genotypes_file_unions_with_training(self, mock_predict, mock_gm, tmp_pa
assert "A1B" in mock_predict["genotypes"]
assert "C2D" in mock_predict["genotypes"]
- def test_genotypes_file_union_preserves_order(self, mock_predict, mock_gm, tmp_path):
+ def test_genotypes_file_union_preserves_order(self, mock_predict, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["C2D"])
predict_growth("cfg.yaml", "post.h5",
@@ -107,7 +107,7 @@ def test_genotypes_file_union_preserves_order(self, mock_predict, mock_gm, tmp_p
idx_c2d = mock_predict["genotypes"].index("C2D")
assert idx_c2d > 0
- def test_concs_file_unions_with_training(self, mock_predict, mock_gm, tmp_path):
+ def test_concs_file_unions_with_training(self, mock_predict, mock_orchestrator, tmp_path):
cf = str(tmp_path / "concs.txt")
_write_lines(cf, [5.0]) # novel concentration
predict_growth("cfg.yaml", "post.h5",
@@ -117,7 +117,7 @@ def test_concs_file_unions_with_training(self, mock_predict, mock_gm, tmp_path):
assert 1.0 in mock_predict["titrant_conc"]
assert 5.0 in mock_predict["titrant_conc"]
- def test_duplicate_concs_not_repeated(self, mock_predict, mock_gm, tmp_path):
+ def test_duplicate_concs_not_repeated(self, mock_predict, mock_orchestrator, tmp_path):
cf = str(tmp_path / "concs.txt")
_write_lines(cf, [1.0]) # already in training
predict_growth("cfg.yaml", "post.h5",
@@ -126,7 +126,7 @@ def test_duplicate_concs_not_repeated(self, mock_predict, mock_gm, tmp_path):
concs = mock_predict["titrant_conc"]
assert concs.count(1.0) == 1
- def test_duplicate_genotypes_not_repeated(self, mock_predict, mock_gm, tmp_path):
+ def test_duplicate_genotypes_not_repeated(self, mock_predict, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["wt"]) # already in training
predict_growth("cfg.yaml", "post.h5",
@@ -176,7 +176,7 @@ def test_only_files_no_file_falls_through_to_none(self, mock_predict, tmp_path):
class TestPredictGrowthTitrantNamesFilter:
- def test_titrant_names_file_filters_output_rows(self, mock_predict, mock_gm, tmp_path):
+ def test_titrant_names_file_filters_output_rows(self, mock_predict, mock_orchestrator, tmp_path):
nf = str(tmp_path / "names.txt")
_write_lines(nf, ["IPTG"])
out = str(tmp_path / "out")
@@ -203,13 +203,13 @@ def test_titrant_names_file_does_not_affect_predict_call(self, mock_predict, tmp
class TestPredictGrowthCheckpointInput:
- def _make_fixtures(self, mock_gm, resolved_path="resolved.h5"):
+ def _make_fixtures(self, mock_orchestrator, resolved_path="resolved.h5"):
"""Return patch stack that intercepts resolve_param_file."""
def fake_predict(**kwargs):
- genotypes = kwargs.get("genotypes") or mock_gm.growth_df["genotype"].unique().tolist()
- concs = kwargs.get("titrant_conc") or mock_gm.growth_df["titrant_conc"].unique().tolist()
+ genotypes = kwargs.get("genotypes") or mock_orchestrator.growth_df["genotype"].unique().tolist()
+ concs = kwargs.get("titrant_conc") or mock_orchestrator.growth_df["titrant_conc"].unique().tolist()
rows = [{"genotype": g, "titrant_name": "IPTG", "titrant_conc": c,
- "median": 10.0}
+ "q0.5": 10.0}
for g in genotypes for c in concs]
return pd.DataFrame(rows)
@@ -217,7 +217,7 @@ def fake_predict(**kwargs):
patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.read_configuration",
- return_value=(mock_gm, {}),
+ return_value=(mock_orchestrator, {}),
),
patch(
"tfscreen.tfmodel.scripts"
@@ -226,15 +226,15 @@ def fake_predict(**kwargs):
),
]
- def test_pkl_param_file_calls_resolve(self, mock_gm, tmp_path):
+ def test_pkl_param_file_calls_resolve(self, mock_orchestrator, tmp_path):
"""resolve_param_file is called when param_file ends with .pkl."""
resolve_calls = []
- def fake_resolve(pf, gm, op):
+ def fake_resolve(pf, orchestrator, op):
resolve_calls.append(pf)
return "resolved.h5"
- patches = self._make_fixtures(mock_gm)
+ patches = self._make_fixtures(mock_orchestrator)
with patches[0], patches[1], patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.resolve_param_file",
@@ -245,15 +245,15 @@ def fake_resolve(pf, gm, op):
assert resolve_calls == ["myrun_checkpoint.pkl"]
- def test_h5_param_file_calls_resolve(self, mock_gm, tmp_path):
+ def test_h5_param_file_calls_resolve(self, mock_orchestrator, tmp_path):
"""resolve_param_file is called for .h5 files too (pass-through)."""
resolve_calls = []
- def fake_resolve(pf, gm, op):
+ def fake_resolve(pf, orchestrator, op):
resolve_calls.append(pf)
return pf
- patches = self._make_fixtures(mock_gm)
+ patches = self._make_fixtures(mock_orchestrator)
with patches[0], patches[1], patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.resolve_param_file",
@@ -264,16 +264,16 @@ def fake_resolve(pf, gm, op):
assert resolve_calls == ["post.h5"]
- def test_resolved_path_passed_to_predict(self, mock_gm, tmp_path):
+ def test_resolved_path_passed_to_predict(self, mock_orchestrator, tmp_path):
"""The path returned by resolve_param_file is what predict receives."""
predict_calls = {}
def fake_predict(**kwargs):
predict_calls["param_posteriors"] = kwargs.get("param_posteriors")
return pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [10.0]})
+ "titrant_conc": [0.0], "q0.5": [10.0]})
- patches = self._make_fixtures(mock_gm)
+ patches = self._make_fixtures(mock_orchestrator)
with patches[0], patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.predict",
@@ -288,16 +288,16 @@ def fake_predict(**kwargs):
assert predict_calls["param_posteriors"] == "resolved_map.h5"
- def test_pkl_passes_point_est_q_to_get(self, mock_gm, tmp_path):
- """q_to_get={"point_est": 0.5} is passed to predict when param_file is .pkl."""
+ def test_pkl_passes_point_est_q_to_get(self, mock_orchestrator, tmp_path):
+ """q_to_get=[0.5] is passed to predict when param_file is .pkl."""
predict_calls = {}
def fake_predict(**kwargs):
predict_calls["q_to_get"] = kwargs.get("q_to_get")
return pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "point_est": [10.0]})
+ "titrant_conc": [0.0], "q0.5": [10.0]})
- patches = self._make_fixtures(mock_gm)
+ patches = self._make_fixtures(mock_orchestrator)
with patches[0], patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.predict",
@@ -310,18 +310,18 @@ def fake_predict(**kwargs):
predict_growth("cfg.yaml", "run_checkpoint.pkl",
out_prefix=str(tmp_path / "out"))
- assert predict_calls["q_to_get"] == {"point_est": 0.5}
+ assert predict_calls["q_to_get"] == [0.5]
- def test_h5_passes_none_q_to_get(self, mock_gm, tmp_path):
+ def test_h5_passes_none_q_to_get(self, mock_orchestrator, tmp_path):
"""q_to_get=None is passed to predict when param_file is .h5."""
predict_calls = {}
def fake_predict(**kwargs):
predict_calls["q_to_get"] = kwargs.get("q_to_get")
return pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [10.0]})
+ "titrant_conc": [0.0], "q0.5": [10.0]})
- patches = self._make_fixtures(mock_gm)
+ patches = self._make_fixtures(mock_orchestrator)
with patches[0], patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.predict",
@@ -329,7 +329,7 @@ def fake_predict(**kwargs):
), patch(
"tfscreen.tfmodel.scripts"
".predict_growth_cli.resolve_param_file",
- side_effect=lambda pf, gm, op: pf,
+ side_effect=lambda pf, orchestrator, op: pf,
):
predict_growth("cfg.yaml", "post.h5",
out_prefix=str(tmp_path / "out"))
diff --git a/tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py b/tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py
index 30ec0334..e140e0a3 100644
--- a/tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_predict_theta_cli.py
@@ -28,16 +28,16 @@ def _make_tm_df(genotypes, titrant_names, titrant_concs):
@pytest.fixture
-def mock_gm():
- gm = MagicMock()
- gm._theta = "hill_geno"
- gm.growth_tm.df = _make_tm_df(
+def mock_orchestrator():
+ orchestrator = MagicMock()
+ orchestrator._theta = "hill_geno"
+ orchestrator.growth_tm.df = _make_tm_df(
["wt", "A1B"],
["IPTG", "IPTG"],
[0.0, 1.0],
)
- gm.training_tm.df = gm.growth_tm.df
- return gm
+ orchestrator.training_tm.df = orchestrator.growth_tm.df
+ return orchestrator
def _fake_extract(model, posteriors, **kwargs):
@@ -48,21 +48,21 @@ def _fake_extract(model, posteriors, **kwargs):
else:
pairs = [("IPTG", 0.0), ("IPTG", 1.0)]
target = kwargs.get("target_genotypes") or ["wt", "A1B"]
- rows = [{"genotype": g, "titrant_name": tn, "titrant_conc": tc, "median": 0.5}
+ rows = [{"genotype": g, "titrant_name": tn, "titrant_conc": tc, "q0.5": 0.5}
for g in target for tn, tc in pairs]
return pd.DataFrame(rows)
@pytest.fixture
-def mock_extract(mock_gm):
+def mock_extract(mock_orchestrator):
with patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.read_configuration",
- return_value=(mock_gm, {}),
+ return_value=(mock_orchestrator, {}),
), patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.resolve_param_file",
- side_effect=lambda pf, gm, op: pf, # pass-through
+ side_effect=lambda pf, orchestrator, op: pf, # pass-through
), patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.extract_theta_curves",
@@ -94,7 +94,7 @@ def test_mismatched_titrant_files_raises(mock_extract, tmp_path):
class TestPredictThetaDefaults:
- def test_no_files_uses_all_training_genotypes(self, mock_extract, mock_gm, tmp_path):
+ def test_no_files_uses_all_training_genotypes(self, mock_extract, mock_orchestrator, tmp_path):
out = str(tmp_path / "out")
with patch(
"tfscreen.tfmodel.scripts"
@@ -102,7 +102,7 @@ def test_no_files_uses_all_training_genotypes(self, mock_extract, mock_gm, tmp_p
) as mock_curves:
mock_curves.return_value = pd.DataFrame(
{"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]}
+ "titrant_conc": [0.0], "q0.5": [0.5]}
)
predict_theta("cfg.yaml", "post.h5", out_prefix=out)
assert mock_curves.called
@@ -114,7 +114,7 @@ def test_no_files_uses_all_training_genotypes(self, mock_extract, mock_gm, tmp_p
class TestPredictThetaUnion:
- def test_genotypes_file_unions_with_training(self, mock_extract, mock_gm, tmp_path):
+ def test_genotypes_file_unions_with_training(self, mock_extract, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["C2D"]) # out-of-training
out = str(tmp_path / "out")
@@ -126,7 +126,7 @@ def test_genotypes_file_unions_with_training(self, mock_extract, mock_gm, tmp_pa
{"genotype": ["wt", "A1B", "C2D"],
"titrant_name": ["IPTG"] * 3,
"titrant_conc": [0.0] * 3,
- "median": [0.5] * 3}
+ "q0.5": [0.5] * 3}
)
predict_theta("cfg.yaml", "post.h5",
genotypes_file=gf, out_prefix=out)
@@ -135,7 +135,7 @@ def test_genotypes_file_unions_with_training(self, mock_extract, mock_gm, tmp_pa
assert "A1B" in target
assert "C2D" in target
- def test_titrant_files_union_with_training_grid(self, mock_extract, mock_gm, tmp_path):
+ def test_titrant_files_union_with_training_grid(self, mock_extract, mock_orchestrator, tmp_path):
nf = str(tmp_path / "names.txt")
cf = str(tmp_path / "concs.txt")
_write_lines(nf, ["IPTG"])
@@ -147,7 +147,7 @@ def test_titrant_files_union_with_training_grid(self, mock_extract, mock_gm, tmp
) as mock_curves:
mock_curves.return_value = pd.DataFrame(
{"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]}
+ "titrant_conc": [0.0], "q0.5": [0.5]}
)
predict_theta("cfg.yaml", "post.h5",
titrant_names_file=nf,
@@ -158,7 +158,7 @@ def test_titrant_files_union_with_training_grid(self, mock_extract, mock_gm, tmp
assert 1.0 in titrant_df["titrant_conc"].values
assert 5.0 in titrant_df["titrant_conc"].values
- def test_single_titrant_name_broadcast_across_concs(self, mock_extract, mock_gm, tmp_path):
+ def test_single_titrant_name_broadcast_across_concs(self, mock_extract, mock_orchestrator, tmp_path):
nf = str(tmp_path / "names.txt")
cf = str(tmp_path / "concs.txt")
_write_lines(nf, ["IPTG"])
@@ -170,7 +170,7 @@ def test_single_titrant_name_broadcast_across_concs(self, mock_extract, mock_gm,
) as mock_curves:
mock_curves.return_value = pd.DataFrame(
{"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]}
+ "titrant_conc": [0.0], "q0.5": [0.5]}
)
predict_theta("cfg.yaml", "post.h5",
titrant_names_file=nf,
@@ -180,7 +180,7 @@ def test_single_titrant_name_broadcast_across_concs(self, mock_extract, mock_gm,
assert list(titrant_df["titrant_name"]) == ["IPTG", "IPTG", "IPTG", "IPTG"]
assert set(titrant_df["titrant_conc"]) >= {0.0, 10.0, 100.0}
- def test_mismatched_titrant_file_lengths_raises(self, mock_extract, mock_gm, tmp_path):
+ def test_mismatched_titrant_file_lengths_raises(self, mock_extract, mock_orchestrator, tmp_path):
nf = str(tmp_path / "names.txt")
cf = str(tmp_path / "concs.txt")
_write_lines(nf, ["IPTG", "ATC"])
@@ -191,7 +191,7 @@ def test_mismatched_titrant_file_lengths_raises(self, mock_extract, mock_gm, tmp
titrant_concs_file=cf,
out_prefix=str(tmp_path / "out"))
- def test_duplicate_titrant_pairs_not_repeated(self, mock_extract, mock_gm, tmp_path):
+ def test_duplicate_titrant_pairs_not_repeated(self, mock_extract, mock_orchestrator, tmp_path):
nf = str(tmp_path / "names.txt")
cf = str(tmp_path / "concs.txt")
_write_lines(nf, ["IPTG"])
@@ -203,7 +203,7 @@ def test_duplicate_titrant_pairs_not_repeated(self, mock_extract, mock_gm, tmp_p
) as mock_curves:
mock_curves.return_value = pd.DataFrame(
{"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]}
+ "titrant_conc": [0.0], "q0.5": [0.5]}
)
predict_theta("cfg.yaml", "post.h5",
titrant_names_file=nf,
@@ -223,7 +223,7 @@ def test_duplicate_titrant_pairs_not_repeated(self, mock_extract, mock_gm, tmp_p
class TestPredictThetaOnlyFiles:
- def test_only_files_restricts_genotypes_to_file(self, mock_extract, mock_gm, tmp_path):
+ def test_only_files_restricts_genotypes_to_file(self, mock_extract, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["A1B"])
out = str(tmp_path / "out")
@@ -233,7 +233,7 @@ def test_only_files_restricts_genotypes_to_file(self, mock_extract, mock_gm, tmp
) as mock_curves:
mock_curves.return_value = pd.DataFrame(
{"genotype": ["A1B"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]}
+ "titrant_conc": [0.0], "q0.5": [0.5]}
)
predict_theta("cfg.yaml", "post.h5",
genotypes_file=gf,
@@ -244,7 +244,7 @@ def test_only_files_restricts_genotypes_to_file(self, mock_extract, mock_gm, tmp
df = pd.read_csv(f"{out}.csv")
assert set(df["genotype"].unique()) <= {"A1B"}
- def test_only_files_restricts_titrant_grid_to_file(self, mock_extract, mock_gm, tmp_path):
+ def test_only_files_restricts_titrant_grid_to_file(self, mock_extract, mock_orchestrator, tmp_path):
nf = str(tmp_path / "names.txt")
cf = str(tmp_path / "concs.txt")
_write_lines(nf, ["IPTG"])
@@ -256,7 +256,7 @@ def test_only_files_restricts_titrant_grid_to_file(self, mock_extract, mock_gm,
) as mock_curves:
mock_curves.return_value = pd.DataFrame(
{"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [5.0], "median": [0.5]}
+ "titrant_conc": [5.0], "q0.5": [0.5]}
)
predict_theta("cfg.yaml", "post.h5",
titrant_names_file=nf,
@@ -275,19 +275,19 @@ def test_only_files_restricts_titrant_grid_to_file(self, mock_extract, mock_gm,
class TestPredictThetaCheckpointInput:
- def _base_patches(self, mock_gm):
+ def _base_patches(self, mock_orchestrator):
return [
patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.read_configuration",
- return_value=(mock_gm, {}),
+ return_value=(mock_orchestrator, {}),
),
patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.extract_theta_curves",
return_value=pd.DataFrame({
"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5],
+ "titrant_conc": [0.0], "q0.5": [0.5],
}),
),
patch(
@@ -295,20 +295,20 @@ def _base_patches(self, mock_gm):
".predict_theta_cli.extract_theta_unmeasured",
return_value=pd.DataFrame({
"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5],
+ "titrant_conc": [0.0], "q0.5": [0.5],
}),
),
]
- def test_pkl_calls_resolve_param_file(self, mock_gm, tmp_path):
+ def test_pkl_calls_resolve_param_file(self, mock_orchestrator, tmp_path):
"""resolve_param_file is invoked when param_file ends with .pkl."""
resolve_calls = []
- def fake_resolve(pf, gm, op):
+ def fake_resolve(pf, orchestrator, op):
resolve_calls.append(pf)
return "resolved.h5"
- p = self._base_patches(mock_gm)
+ p = self._base_patches(mock_orchestrator)
with p[0], p[1], p[2], patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.resolve_param_file",
@@ -319,7 +319,7 @@ def fake_resolve(pf, gm, op):
assert resolve_calls == ["run_checkpoint.pkl"]
- def test_resolved_path_passed_to_extract(self, mock_gm, tmp_path):
+ def test_resolved_path_passed_to_extract(self, mock_orchestrator, tmp_path):
"""The path returned by resolve_param_file reaches extract_theta_curves."""
extract_calls = {}
@@ -327,10 +327,10 @@ def fake_curves(**kwargs):
extract_calls["posteriors"] = kwargs.get("posteriors")
return pd.DataFrame({
"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5],
+ "titrant_conc": [0.0], "q0.5": [0.5],
})
- p = self._base_patches(mock_gm)
+ p = self._base_patches(mock_orchestrator)
with p[0], patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.extract_theta_curves",
@@ -353,7 +353,7 @@ def fake_curves(**kwargs):
class TestGenotypeBatchSize:
"""genotype_batch_size is forwarded to extract_theta_unmeasured."""
- def test_custom_batch_size_forwarded(self, mock_extract, mock_gm, tmp_path):
+ def test_custom_batch_size_forwarded(self, mock_extract, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["C2D"]) # out-of-training → unmeasured path
out = str(tmp_path / "out")
@@ -365,7 +365,7 @@ def test_custom_batch_size_forwarded(self, mock_extract, mock_gm, tmp_path):
"genotype": ["wt", "A1B", "C2D"],
"titrant_name": ["IPTG"] * 3,
"titrant_conc": [0.0] * 3,
- "median": [0.5] * 3,
+ "q0.5": [0.5] * 3,
})
predict_theta("cfg.yaml", "post.h5",
genotypes_file=gf,
@@ -373,7 +373,7 @@ def test_custom_batch_size_forwarded(self, mock_extract, mock_gm, tmp_path):
out_prefix=out)
assert mock_unmeas.call_args.kwargs["genotype_batch_size"] == 42
- def test_default_batch_size_is_2000(self, mock_extract, mock_gm, tmp_path):
+ def test_default_batch_size_is_2000(self, mock_extract, mock_orchestrator, tmp_path):
gf = str(tmp_path / "genos.txt")
_write_lines(gf, ["C2D"])
out = str(tmp_path / "out")
@@ -385,7 +385,7 @@ def test_default_batch_size_is_2000(self, mock_extract, mock_gm, tmp_path):
"genotype": ["wt", "A1B", "C2D"],
"titrant_name": ["IPTG"] * 3,
"titrant_conc": [0.0] * 3,
- "median": [0.5] * 3,
+ "q0.5": [0.5] * 3,
})
predict_theta("cfg.yaml", "post.h5",
genotypes_file=gf,
@@ -399,19 +399,19 @@ def test_default_batch_size_is_2000(self, mock_extract, mock_gm, tmp_path):
class TestPredictThetaQToGet:
- def test_pkl_passes_point_est_q_to_get(self, mock_gm, tmp_path):
- """q_to_get={"point_est": 0.5} is passed to extract_theta_curves for .pkl input."""
+ def test_pkl_passes_point_est_q_to_get(self, mock_orchestrator, tmp_path):
+ """q_to_get=[0.5] is passed to extract_theta_curves for .pkl input."""
extract_calls = {}
def fake_curves(**kwargs):
extract_calls["q_to_get"] = kwargs.get("q_to_get")
return pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "point_est": [0.5]})
+ "titrant_conc": [0.0], "q0.5": [0.5]})
with patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.read_configuration",
- return_value=(mock_gm, {}),
+ return_value=(mock_orchestrator, {}),
), patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.resolve_param_file",
@@ -424,30 +424,30 @@ def fake_curves(**kwargs):
"tfscreen.tfmodel.scripts"
".predict_theta_cli.extract_theta_unmeasured",
return_value=pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "point_est": [0.5]}),
+ "titrant_conc": [0.0], "q0.5": [0.5]}),
):
predict_theta("cfg.yaml", "run_checkpoint.pkl",
out_prefix=str(tmp_path / "out"))
- assert extract_calls["q_to_get"] == {"point_est": 0.5}
+ assert extract_calls["q_to_get"] == [0.5]
- def test_h5_passes_none_q_to_get(self, mock_gm, tmp_path):
+ def test_h5_passes_none_q_to_get(self, mock_orchestrator, tmp_path):
"""q_to_get=None is passed to extract_theta_curves for .h5 input."""
extract_calls = {}
def fake_curves(**kwargs):
extract_calls["q_to_get"] = kwargs.get("q_to_get")
return pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]})
+ "titrant_conc": [0.0], "q0.5": [0.5]})
with patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.read_configuration",
- return_value=(mock_gm, {}),
+ return_value=(mock_orchestrator, {}),
), patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.resolve_param_file",
- side_effect=lambda pf, gm, op: pf,
+ side_effect=lambda pf, orchestrator, op: pf,
), patch(
"tfscreen.tfmodel.scripts"
".predict_theta_cli.extract_theta_curves",
@@ -456,7 +456,7 @@ def fake_curves(**kwargs):
"tfscreen.tfmodel.scripts"
".predict_theta_cli.extract_theta_unmeasured",
return_value=pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"],
- "titrant_conc": [0.0], "median": [0.5]}),
+ "titrant_conc": [0.0], "q0.5": [0.5]}),
):
predict_theta("cfg.yaml", "post.h5",
out_prefix=str(tmp_path / "out"))
diff --git a/tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py b/tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py
new file mode 100644
index 00000000..05c5c3fc
--- /dev/null
+++ b/tests/tfscreen/tfmodel/scripts/test_prefit_calibration_cli.py
@@ -0,0 +1,1345 @@
+"""
+Tests for prefit_calibration_cli.py.
+
+The pre-fit script builds an in-process calibration ModelOrchestrator,
+runs MAP, and then writes in-place updates (with .bak backups) into the
+production priors and guesses CSVs. These tests exercise the helper
+functions directly and mock the heavy machinery (read_configuration,
+ModelOrchestrator, RunInference) when testing the orchestration in
+run_prefit_calibration / main.
+"""
+import dataclasses
+import os
+import shutil
+import sys
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pandas as pd
+import pytest
+import yaml
+import jax
+import jax.numpy as jnp
+
+import flax.struct as fstruct
+
+from tfscreen.tfmodel.scripts.prefit_calibration_cli import (
+ _apply_guesses_updates,
+ _apply_priors_updates,
+ _build_calibration_model,
+ _build_csv_updates,
+ _build_hessian_scale_updates,
+ _compute_theta_values,
+ _csv_row_name,
+ _identify_field_mapping,
+ _inject_calibration_priors,
+ _intersect_data,
+ _resolve_csv_paths,
+ main,
+ run_prefit_calibration,
+ _CALIBRATION_OVERRIDES,
+ _DEFAULT_K_SCALE_CEILING,
+ _DEFAULT_K_SCALE_FLOOR,
+ _DEFAULT_M_SCALE_CEILING,
+ _DEFAULT_M_SCALE_FLOOR,
+ _PINNED_COMPONENTS,
+)
+
+
+# ---------------------------------------------------------------------------
+# Lightweight stand-ins for the flax pytree dataclasses used by ModelPriors.
+#
+# The real component-prior dataclasses are flax.struct.dataclass instances
+# that expose a .replace() method. We define a few minimal substitutes
+# here so the unit tests don't need to instantiate a full ModelOrchestrator.
+# ---------------------------------------------------------------------------
+
+@fstruct.dataclass
+class _FakeCondGrowth:
+ growth_k_hyper_loc: float = 0.0
+ k_hyper_loc_scale: float = 1.0
+ growth_k_hyper_scale_loc: float = 0.5
+ growth_m_hyper_loc: float = 0.0
+ m_hyper_loc_scale: float = 1.0
+ growth_m_hyper_scale_loc: float = 0.5
+ pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
+
+
+@fstruct.dataclass
+class _FakeGrowthTransition:
+ pre_t_hyper_loc: float = 0.0
+ pre_t_hyper_loc_scale: float = 1.0
+ pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
+
+
+@fstruct.dataclass
+class _FakeGTNoPinned:
+ pre_t_hyper_loc: float = 0.0
+ pre_t_hyper_loc_scale: float = 1.0
+ # No pinned field; mirrors the "instant" growth_transition variant.
+
+
+@fstruct.dataclass
+class _FakeActivity:
+ hyper_loc_loc: float = 0.0
+ hyper_scale_loc: float = 1.0
+ pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
+
+
+@fstruct.dataclass
+class _FakeDkGeno:
+ hyper_loc_loc: float = 0.0
+ hyper_scale_loc: float = 1.0
+ hyper_shift_loc: float = 0.5
+ pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
+
+
+@fstruct.dataclass
+class _FakeLnCfu0:
+ # Array fields (one element per library class) matching the real ModelPriors.
+ ln_cfu0_hyper_loc_locs: object = None
+ ln_cfu0_hyper_scale_locs: object = None
+ pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
+
+
+@fstruct.dataclass
+class _FakeTheta:
+ theta_values: object = None # array placeholder
+
+
+@fstruct.dataclass
+class _FakeGrowth:
+ condition_growth: _FakeCondGrowth
+ growth_transition: _FakeGrowthTransition
+ activity: _FakeActivity
+ dk_geno: _FakeDkGeno
+ ln_cfu0: _FakeLnCfu0
+
+
+@fstruct.dataclass
+class _FakePriors:
+ growth: _FakeGrowth
+ theta: _FakeTheta
+
+
+def _make_fake_priors(gt_has_pinned=True):
+ gt = _FakeGrowthTransition() if gt_has_pinned else _FakeGTNoPinned()
+ return _FakePriors(
+ growth=_FakeGrowth(
+ condition_growth=_FakeCondGrowth(),
+ growth_transition=gt,
+ activity=_FakeActivity(),
+ dk_geno=_FakeDkGeno(),
+ ln_cfu0=_FakeLnCfu0(
+ ln_cfu0_hyper_loc_locs=jnp.array([5.0]),
+ ln_cfu0_hyper_scale_locs=jnp.array([1.0]),
+ ),
+ ),
+ theta=_FakeTheta(),
+ )
+
+
+# ---------------------------------------------------------------------------
+# _intersect_data
+# ---------------------------------------------------------------------------
+
+class TestIntersectData:
+
+ def _basic_growth(self):
+ return pd.DataFrame({
+ "genotype": ["wt", "wt", "A1T", "G2P"],
+ "titrant_name": ["IPTG"] * 4,
+ "titrant_conc": [0.0, 1.0, 0.0, 0.0],
+ "ln_cfu": [1.0, 2.0, 3.0, 4.0],
+ })
+
+ def _basic_binding(self):
+ return pd.DataFrame({
+ "genotype": ["wt", "wt", "A1T", "M3L"],
+ "titrant_name": ["IPTG"] * 4,
+ "titrant_conc": [0.0, 1.0, 0.0, 0.0],
+ "theta_obs": [0.1, 0.5, 0.2, 0.3],
+ "theta_std": [0.05, 0.05, 0.05, 0.05],
+ })
+
+ def test_intersection_keeps_only_shared_rows(self):
+ g, b = _intersect_data(self._basic_growth(), self._basic_binding())
+ # G2P is only in growth, M3L is only in binding; both must drop.
+ assert set(g["genotype"].unique()) == {"wt", "A1T"}
+ assert set(b["genotype"].unique()) == {"wt", "A1T"}
+
+ def test_intersection_returns_copies_not_views(self):
+ gdf = self._basic_growth()
+ bdf = self._basic_binding()
+ g, _ = _intersect_data(gdf, bdf)
+ # Mutating the returned frame must not alter the original.
+ g.loc[g.index[0], "ln_cfu"] = -999.0
+ assert gdf.loc[0, "ln_cfu"] != -999.0
+
+ def test_intersection_preserves_per_titrant_resolution(self):
+ # (wt, IPTG, 1.0) and (wt, IPTG, 0.0) must each be evaluated
+ # separately — we cannot just intersect on genotype.
+ growth = pd.DataFrame({
+ "genotype": ["wt", "wt"],
+ "titrant_name": ["IPTG", "IPTG"],
+ "titrant_conc": [0.0, 1.0],
+ "ln_cfu": [1.0, 2.0],
+ })
+ binding = pd.DataFrame({
+ "genotype": ["wt"],
+ "titrant_name": ["IPTG"],
+ "titrant_conc": [1.0],
+ "theta_obs": [0.5],
+ "theta_std": [0.05],
+ })
+ g, b = _intersect_data(growth, binding)
+ assert len(g) == 1
+ assert g["titrant_conc"].iloc[0] == 1.0
+
+ def test_empty_intersection_raises(self):
+ growth = pd.DataFrame({
+ "genotype": ["wt"],
+ "titrant_name": ["IPTG"],
+ "titrant_conc": [0.0],
+ "ln_cfu": [1.0],
+ })
+ binding = pd.DataFrame({
+ "genotype": ["A1T"],
+ "titrant_name": ["IPTG"],
+ "titrant_conc": [0.0],
+ "theta_obs": [0.5],
+ "theta_std": [0.05],
+ })
+ with pytest.raises(ValueError, match="intersection is empty"):
+ _intersect_data(growth, binding)
+
+ def test_missing_column_in_growth_raises(self):
+ growth = pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"]})
+ binding = self._basic_binding()
+ with pytest.raises(ValueError, match="growth_df is missing"):
+ _intersect_data(growth, binding)
+
+ def test_missing_column_in_binding_raises(self):
+ growth = self._basic_growth()
+ binding = pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"]})
+ with pytest.raises(ValueError, match="binding_df is missing"):
+ _intersect_data(growth, binding)
+
+
+# ---------------------------------------------------------------------------
+# _compute_theta_values
+# ---------------------------------------------------------------------------
+
+class TestComputeThetaValues:
+
+ def _make_orchestrator_cal(self, titrant_names, titrant_concs, genotypes):
+ """Return a MagicMock with the binding_tm shape used by the helper."""
+ orchestrator = MagicMock()
+ # Order matters; the helper reads names, concs, and genotypes by index.
+ orchestrator.binding_tm.tensor_dim_names = [
+ "titrant_name", "titrant_conc", "genotype"
+ ]
+ orchestrator.binding_tm.tensor_dim_labels = [
+ list(titrant_names),
+ list(titrant_concs),
+ list(genotypes),
+ ]
+ return orchestrator
+
+ def test_per_genotype_theta_extracted_correctly(self):
+ """Each genotype must get its own theta value, not the cross-genotype average."""
+ orchestrator = self._make_orchestrator_cal(["IPTG"], [1.0], ["wt", "mut"])
+ # wt has theta=0.8, mut has theta=0.2; output must preserve the distinction.
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG", "IPTG"],
+ "titrant_conc": [1.0, 1.0],
+ "genotype": ["wt", "mut"],
+ "theta_obs": [0.8, 0.2],
+ "theta_std": [0.01, 0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (1, 1, 2)
+ assert theta[0, 0, 0] > 0.75 # wt ≈ 0.8
+ assert theta[0, 0, 1] < 0.25 # mut ≈ 0.2
+
+ def test_inverse_variance_weighting_within_genotype(self):
+ """Multiple rows for the same genotype/cell must be IVW-averaged."""
+ orchestrator = self._make_orchestrator_cal(["IPTG"], [1.0], ["wt"])
+ # Two replicate measurements: theta=0.2 (noisy) and theta=0.8 (precise).
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG", "IPTG"],
+ "titrant_conc": [1.0, 1.0],
+ "genotype": ["wt", "wt"],
+ "theta_obs": [0.2, 0.8],
+ "theta_std": [0.1, 0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (1, 1, 1)
+ # Inverse-variance weighting heavily favours the precise observation.
+ assert theta[0, 0, 0] > 0.7
+
+ def test_falls_back_to_plain_mean_when_all_stds_zero(self):
+ orchestrator = self._make_orchestrator_cal(["IPTG"], [1.0], ["wt"])
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG", "IPTG"],
+ "titrant_conc": [1.0, 1.0],
+ "genotype": ["wt", "wt"],
+ "theta_obs": [0.3, 0.7],
+ "theta_std": [0.0, 0.0], # invalid weights → fallback to mean
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta[0, 0, 0] == pytest.approx(0.5)
+
+ def test_unobserved_genotype_cell_defaults_to_midpoint(self):
+ """A genotype present in labels but absent from binding_df gets 0.5."""
+ orchestrator = self._make_orchestrator_cal(["IPTG"], [1.0], ["wt", "missing"])
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG"],
+ "titrant_conc": [1.0],
+ "genotype": ["wt"],
+ "theta_obs": [0.9],
+ "theta_std": [0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (1, 1, 2)
+ assert theta[0, 0, 0] > 0.85 # wt observed
+ assert theta[0, 0, 1] == pytest.approx(0.5) # "missing" → midpoint
+
+ def test_unobserved_concentration_cell_defaults_to_midpoint(self):
+ """A concentration present in labels but absent from binding_df gets 0.5."""
+ orchestrator = self._make_orchestrator_cal(["IPTG"], [0.0, 1.0], ["wt"])
+ # Only the IPTG=1.0 cell is observed.
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG"],
+ "titrant_conc": [1.0],
+ "genotype": ["wt"],
+ "theta_obs": [0.9],
+ "theta_std": [0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (1, 2, 1)
+ assert theta[0, 0, 0] == pytest.approx(0.5) # conc=0.0 not observed
+ assert theta[0, 1, 0] > 0.85 # conc=1.0 observed
+
+ def test_output_shape(self):
+ """Output shape must be (n_titrant_name, n_titrant_conc, n_genotype)."""
+ orchestrator = self._make_orchestrator_cal(
+ ["IPTG", "arabinose"], [0.0, 0.01, 1.0], ["wt", "A1T", "G2P"]
+ )
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG"] * 3,
+ "titrant_conc": [0.0, 0.01, 1.0],
+ "genotype": ["wt", "A1T", "G2P"],
+ "theta_obs": [0.9, 0.5, 0.1],
+ "theta_std": [0.01, 0.01, 0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (2, 3, 3)
+
+ def test_clip_to_open_interval(self):
+ orchestrator = self._make_orchestrator_cal(["IPTG"], [0.0, 1.0], ["wt"])
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG", "IPTG"],
+ "titrant_conc": [0.0, 1.0],
+ "genotype": ["wt", "wt"],
+ "theta_obs": [0.0, 1.0],
+ "theta_std": [0.01, 0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (1, 2, 1)
+ # Zero and one would blow up the downstream logit; expect clipping.
+ assert theta[0, 0, 0] > 0.0
+ assert theta[0, 1, 0] < 1.0
+
+ def test_genotype_ordering_matches_labels(self):
+ """Values must be placed at the index matching the genotype label order."""
+ orchestrator = self._make_orchestrator_cal(
+ ["IPTG"], [1.0], ["wt", "A1T", "G2P"]
+ )
+ binding_df = pd.DataFrame({
+ "titrant_name": ["IPTG", "IPTG", "IPTG"],
+ "titrant_conc": [1.0, 1.0, 1.0],
+ # Deliberately supply rows in a different order than the labels.
+ "genotype": ["G2P", "wt", "A1T"],
+ "theta_obs": [0.1, 0.9, 0.5],
+ "theta_std": [0.01, 0.01, 0.01],
+ })
+ theta = np.asarray(_compute_theta_values(orchestrator, binding_df))
+ assert theta.shape == (1, 1, 3)
+ # wt is index 0 in labels → theta ≈ 0.9
+ assert theta[0, 0, 0] > 0.85
+ # A1T is index 1 → theta ≈ 0.5
+ assert 0.45 < theta[0, 0, 1] < 0.55
+ # G2P is index 2 → theta ≈ 0.1
+ assert theta[0, 0, 2] < 0.15
+
+
+# ---------------------------------------------------------------------------
+# _csv_row_name / _build_csv_updates
+# ---------------------------------------------------------------------------
+
+class TestCsvRowName:
+
+ def test_basic_row_name(self):
+ assert _csv_row_name("condition_growth", "growth_k_hyper_loc") == \
+ "growth.condition_growth.growth_k_hyper_loc"
+
+
+class TestBuildCsvUpdates:
+
+ def test_normal_site_writes_loc_to_priors_and_guesses(self):
+ field_mapping = {
+ "condition_growth_growth_k_hyper": {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": "growth_k_hyper_loc",
+ "scale_field": "k_hyper_loc_scale",
+ "is_array": False,
+ },
+ }
+ params = {"condition_growth_growth_k_hyper_auto_loc": np.float32(2.5)}
+ prior_updates, guess_updates = _build_csv_updates(field_mapping, params)
+ # MAP → loc_field in priors
+ assert prior_updates[
+ "growth.condition_growth.growth_k_hyper_loc"] == pytest.approx(2.5)
+ # scale_field is NOT updated (no Hessian sigma)
+ assert "growth.condition_growth.k_hyper_loc_scale" not in prior_updates
+ # MAP → scalar guess
+ assert guess_updates["condition_growth_growth_k_hyper"] == pytest.approx(2.5)
+
+ def test_halfnormal_site_writes_map_to_scale_field(self):
+ field_mapping = {
+ "condition_growth_k_hyper_scale": {
+ "component": "condition_growth",
+ "dist_class": "HalfNormal",
+ "scale_field": "growth_k_hyper_scale_loc",
+ "is_array": False,
+ },
+ }
+ params = {"condition_growth_k_hyper_scale_auto_loc": np.float32(0.42)}
+ prior_updates, guess_updates = _build_csv_updates(field_mapping, params)
+ assert prior_updates == {
+ "growth.condition_growth.growth_k_hyper_scale_loc": pytest.approx(0.42),
+ }
+ assert guess_updates["condition_growth_k_hyper_scale"] == pytest.approx(0.42)
+
+ def test_skips_sites_without_auto_loc_in_params(self):
+ field_mapping = {
+ "missing_site": {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": "growth_k_hyper_loc",
+ "scale_field": "k_hyper_loc_scale",
+ "is_array": False,
+ },
+ }
+ prior_updates, guess_updates = _build_csv_updates(field_mapping, {})
+ assert prior_updates == {}
+ assert guess_updates == {}
+
+ def test_array_site_writes_locs_to_guesses(self):
+ """Simple-prior array sites write per-condition MAP values to guesses."""
+ field_mapping = {
+ "condition_growth_k": {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": "k_loc",
+ "scale_field": "k_scale",
+ "is_array": True,
+ },
+ }
+ params = {"condition_growth_k_auto_loc": np.array([1.0, 2.0, 3.0])}
+ prior_updates, guess_updates = _build_csv_updates(field_mapping, params)
+ # Array sites do NOT update priors
+ assert prior_updates == {}
+ # Guess locs get the per-condition MAP array
+ assert "condition_growth_k_locs" in guess_updates
+ assert np.allclose(guess_updates["condition_growth_k_locs"],
+ [1.0, 2.0, 3.0])
+
+
+# ---------------------------------------------------------------------------
+# _apply_priors_updates / _apply_guesses_updates
+# ---------------------------------------------------------------------------
+
+class TestApplyPriorsUpdates:
+
+ def _write_csv(self, tmp_path, rows):
+ path = tmp_path / "priors.csv"
+ pd.DataFrame(rows).to_csv(path, index=False)
+ return str(path)
+
+ def test_writes_bak_and_updates_only_matching_rows(self, tmp_path):
+ path = self._write_csv(tmp_path, [
+ {"parameter": "growth.condition_growth.growth_k_hyper_loc",
+ "value": 1.0},
+ {"parameter": "growth.activity.hyper_loc_loc",
+ "value": 0.0},
+ ])
+ updates = {"growth.condition_growth.growth_k_hyper_loc": 9.0}
+ _apply_priors_updates(path, updates)
+
+ bak = pd.read_csv(path + ".bak")
+ new = pd.read_csv(path)
+ # Backup retains the original
+ assert bak.set_index("parameter")["value"][
+ "growth.condition_growth.growth_k_hyper_loc"] == 1.0
+ # Live file has the updated value
+ assert new.set_index("parameter")["value"][
+ "growth.condition_growth.growth_k_hyper_loc"] == 9.0
+ # Untouched rows are preserved verbatim
+ assert new.set_index("parameter")["value"][
+ "growth.activity.hyper_loc_loc"] == 0.0
+
+ def test_no_updates_is_noop(self, tmp_path):
+ path = self._write_csv(tmp_path, [
+ {"parameter": "x", "value": 1.0},
+ ])
+ _apply_priors_updates(path, {})
+ # No .bak, no rewrite.
+ assert not os.path.exists(path + ".bak")
+
+ def test_warns_on_missing_row(self, tmp_path, capsys):
+ path = self._write_csv(tmp_path, [
+ {"parameter": "x", "value": 1.0},
+ ])
+ _apply_priors_updates(path, {"y_does_not_exist": 5.0})
+ err = capsys.readouterr().err
+ assert "no matching row" in err
+ assert "y_does_not_exist" in err
+
+ def test_raises_on_malformed_csv(self, tmp_path):
+ bad = tmp_path / "bad.csv"
+ pd.DataFrame({"foo": [1, 2]}).to_csv(bad, index=False)
+ with pytest.raises(ValueError, match="missing required"):
+ _apply_priors_updates(str(bad), {"x": 1.0})
+
+
+class TestApplyGuessesUpdates:
+
+ def _write_csv(self, tmp_path, rows):
+ path = tmp_path / "guesses.csv"
+ pd.DataFrame(rows).to_csv(path, index=False)
+ return str(path)
+
+ def test_writes_bak_and_updates_only_scalar_rows(self, tmp_path):
+ # Two rows for the same parameter: one scalar (flat_index NaN),
+ # one array entry (flat_index=0). Only the scalar must be
+ # overwritten.
+ path = self._write_csv(tmp_path, [
+ {"parameter": "growth_k_hyper", "value": 1.0,
+ "flat_index": float("nan")},
+ {"parameter": "growth_k_offset", "value": 0.0, "flat_index": 0.0},
+ {"parameter": "growth_k_offset", "value": 0.0, "flat_index": 1.0},
+ ])
+ updates = {"growth_k_hyper": 7.7,
+ "growth_k_offset": 999.0}
+ _apply_guesses_updates(path, updates)
+
+ new = pd.read_csv(path)
+ # Scalar row updated
+ scalar_row = new[(new["parameter"] == "growth_k_hyper")
+ & (new["flat_index"].isna())]
+ assert scalar_row["value"].iloc[0] == pytest.approx(7.7)
+ # Array rows untouched (the update key matched no scalar row)
+ offset_rows = new[new["parameter"] == "growth_k_offset"]
+ assert (offset_rows["value"] == 0.0).all()
+
+ # .bak written
+ assert os.path.exists(path + ".bak")
+
+ def test_no_updates_is_noop(self, tmp_path):
+ path = self._write_csv(tmp_path, [
+ {"parameter": "x", "value": 1.0, "flat_index": float("nan")},
+ ])
+ _apply_guesses_updates(path, {})
+ assert not os.path.exists(path + ".bak")
+
+ def test_updates_array_rows_by_flat_index(self, tmp_path):
+ path = self._write_csv(tmp_path, [
+ {"parameter": "condition_growth_k_locs", "value": 0.0, "flat_index": 0.0},
+ {"parameter": "condition_growth_k_locs", "value": 0.0, "flat_index": 1.0},
+ {"parameter": "condition_growth_k_locs", "value": 0.0, "flat_index": 2.0},
+ {"parameter": "other_param", "value": 9.9, "flat_index": float("nan")},
+ ])
+ _apply_guesses_updates(path, {"condition_growth_k_locs": np.array([1.1, 2.2, 3.3])})
+ new = pd.read_csv(path)
+ k_rows = new[new["parameter"] == "condition_growth_k_locs"].sort_values("flat_index")
+ assert k_rows["value"].tolist() == pytest.approx([1.1, 2.2, 3.3])
+ # Unrelated row preserved
+ assert new[new["parameter"] == "other_param"]["value"].iloc[0] == pytest.approx(9.9)
+ assert os.path.exists(path + ".bak")
+
+ def test_treats_all_rows_scalar_when_no_flat_index_column(self, tmp_path):
+ path = self._write_csv(tmp_path, [
+ {"parameter": "alpha", "value": 1.0},
+ ])
+ _apply_guesses_updates(path, {"alpha": 4.5})
+ new = pd.read_csv(path)
+ assert new.set_index("parameter")["value"]["alpha"] == 4.5
+
+ def test_raises_on_malformed_csv(self, tmp_path):
+ bad = tmp_path / "bad.csv"
+ pd.DataFrame({"foo": [1]}).to_csv(bad, index=False)
+ with pytest.raises(ValueError, match="missing required"):
+ _apply_guesses_updates(str(bad), {"x": 1.0})
+
+
+# ---------------------------------------------------------------------------
+# _resolve_csv_paths
+# ---------------------------------------------------------------------------
+
+class TestResolveCsvPaths:
+
+ def _write_yaml(self, path, body):
+ with open(path, "w") as fh:
+ yaml.dump(body, fh)
+
+ def test_returns_resolved_paths(self, tmp_path):
+ priors = tmp_path / "p.csv"
+ guesses = tmp_path / "g.csv"
+ priors.write_text("parameter,value\n")
+ guesses.write_text("parameter,value\n")
+ cfg = tmp_path / "cfg.yaml"
+ self._write_yaml(cfg, {"priors_file": "p.csv",
+ "guesses_file": "g.csv"})
+ p, g = _resolve_csv_paths(str(cfg))
+ assert os.path.abspath(p) == str(priors)
+ assert os.path.abspath(g) == str(guesses)
+
+ def test_absolute_paths_are_honored(self, tmp_path):
+ priors = tmp_path / "abs_p.csv"
+ guesses = tmp_path / "abs_g.csv"
+ priors.write_text("parameter,value\n")
+ guesses.write_text("parameter,value\n")
+ cfg = tmp_path / "cfg.yaml"
+ self._write_yaml(cfg, {"priors_file": str(priors),
+ "guesses_file": str(guesses)})
+ p, g = _resolve_csv_paths(str(cfg))
+ assert p == str(priors)
+ assert g == str(guesses)
+
+ def test_missing_priors_key_raises(self, tmp_path):
+ cfg = tmp_path / "cfg.yaml"
+ self._write_yaml(cfg, {"guesses_file": "g.csv"})
+ with pytest.raises(ValueError, match="priors_file"):
+ _resolve_csv_paths(str(cfg))
+
+ def test_missing_guesses_key_raises(self, tmp_path):
+ cfg = tmp_path / "cfg.yaml"
+ self._write_yaml(cfg, {"priors_file": "p.csv"})
+ with pytest.raises(ValueError, match="guesses_file"):
+ _resolve_csv_paths(str(cfg))
+
+ def test_missing_priors_file_raises(self, tmp_path):
+ cfg = tmp_path / "cfg.yaml"
+ guesses = tmp_path / "g.csv"
+ guesses.write_text("parameter,value\n")
+ self._write_yaml(cfg, {"priors_file": "p_missing.csv",
+ "guesses_file": "g.csv"})
+ with pytest.raises(FileNotFoundError, match="Priors file not found"):
+ _resolve_csv_paths(str(cfg))
+
+ def test_missing_guesses_file_raises(self, tmp_path):
+ cfg = tmp_path / "cfg.yaml"
+ priors = tmp_path / "p.csv"
+ priors.write_text("parameter,value\n")
+ self._write_yaml(cfg, {"priors_file": "p.csv",
+ "guesses_file": "g_missing.csv"})
+ with pytest.raises(FileNotFoundError, match="Guesses file not found"):
+ _resolve_csv_paths(str(cfg))
+
+
+# ---------------------------------------------------------------------------
+# _build_calibration_model — verify it forwards the right overrides
+# ---------------------------------------------------------------------------
+
+class TestBuildCalibrationModel:
+
+ def test_overrides_replace_production_components(self):
+ orchestrator_prod = MagicMock()
+ orchestrator_prod.settings = {
+ "theta": "categorical_geno", # → simple
+ "activity": "horseshoe_geno", # → hierarchical
+ "dk_geno": "fixed", # → hierarchical
+ "ln_cfu0": "fixed", # → hierarchical
+ "transformation": "logit_norm", # → single
+ "theta_growth_noise": "beta", # → zero
+ "theta_binding_noise": "beta", # → zero
+ "condition_growth": "linear", # passthrough
+ "growth_transition": "instant", # passthrough
+ "batch_size": 7,
+ "spiked_genotypes": ["WT"],
+ }
+
+ with patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.ModelOrchestrator"
+ ) as MockGM:
+ MockGM.return_value = MagicMock()
+ growth_df = pd.DataFrame()
+ binding_df = pd.DataFrame()
+ _build_calibration_model(orchestrator_prod, growth_df, binding_df)
+
+ kwargs = MockGM.call_args.kwargs
+ # Overrides applied
+ for k, v in _CALIBRATION_OVERRIDES.items():
+ assert kwargs[k] == v
+ # Spiked genotypes dropped (calibration only sees the intersection)
+ assert kwargs["spiked_genotypes"] is None
+ # batch_size pulled out of settings and passed positionally
+ assert kwargs["batch_size"] == 7
+ # condition_growth and growth_transition flow through unchanged
+ assert kwargs["condition_growth"] == "linear"
+ assert kwargs["growth_transition"] == "instant"
+ # binding_weight must be 1.0 regardless of production value so the
+ # calibration MAP learns the binding→growth linkage without the
+ # production upweighting drowning the binding signal.
+ assert kwargs["binding_weight"] == 1.0
+
+ def test_binding_weight_reset_from_large_production_value(self):
+ # Reproduces the weighting bug: production YAML stores a large
+ # binding_weight (N_growth / N_binding); the calibration model
+ # must use 1.0 instead.
+ orchestrator_prod = MagicMock()
+ orchestrator_prod.settings = {
+ "theta": "categorical_geno",
+ "activity": "horseshoe_geno",
+ "dk_geno": "fixed",
+ "ln_cfu0": "fixed",
+ "transformation": "logit_norm",
+ "theta_growth_noise": "beta",
+ "theta_binding_noise": "beta",
+ "condition_growth": "linear",
+ "growth_transition": "instant",
+ "batch_size": None,
+ "spiked_genotypes": None,
+ "binding_weight": 250.0, # typical large production value
+ }
+
+ with patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.ModelOrchestrator"
+ ) as MockGM:
+ MockGM.return_value = MagicMock()
+ _build_calibration_model(orchestrator_prod, pd.DataFrame(), pd.DataFrame())
+
+ assert MockGM.call_args.kwargs["binding_weight"] == 1.0
+
+ def test_does_not_mutate_production_settings(self):
+ # Ensure we work on a copy.
+ orchestrator_prod = MagicMock()
+ orchestrator_prod.settings = {"theta": "categorical_geno",
+ "activity": "horseshoe_geno",
+ "dk_geno": "fixed",
+ "ln_cfu0": "fixed",
+ "transformation": "logit_norm",
+ "theta_growth_noise": "beta",
+ "theta_binding_noise": "beta",
+ "spiked_genotypes": None,
+ "batch_size": None,
+ "binding_weight": 150.0}
+ original_settings = dict(orchestrator_prod.settings)
+
+ with patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.ModelOrchestrator"
+ ) as MockGM:
+ MockGM.return_value = MagicMock()
+ _build_calibration_model(orchestrator_prod, pd.DataFrame(), pd.DataFrame())
+
+ assert orchestrator_prod.settings == original_settings
+
+
+# ---------------------------------------------------------------------------
+# _inject_calibration_priors — end-to-end pin / copy behaviour
+# ---------------------------------------------------------------------------
+
+class TestInjectCalibrationPriors:
+
+ def _make_models(self, gt_has_pinned=True):
+ orchestrator_cal = MagicMock()
+ orchestrator_prod = MagicMock()
+ orchestrator_cal._priors_history = []
+
+ cal_priors = _make_fake_priors(gt_has_pinned=gt_has_pinned)
+ prod_priors = _make_fake_priors(gt_has_pinned=True)
+ # Tweak production scalars so we can verify they get copied.
+ prod_priors = prod_priors.replace(
+ growth=prod_priors.growth.replace(
+ condition_growth=prod_priors.growth.condition_growth.replace(
+ growth_k_hyper_loc=2.5,
+ k_hyper_loc_scale=0.4,
+ ),
+ growth_transition=prod_priors.growth.growth_transition.replace(
+ pre_t_hyper_loc=3.5,
+ ),
+ ),
+ )
+ orchestrator_cal.priors = cal_priors
+ orchestrator_prod.priors = prod_priors
+
+ # Capture all writes to _priors so the test can introspect them.
+ def _setter(value):
+ orchestrator_cal._priors_history.append(value)
+ type(orchestrator_cal)._priors = property(
+ lambda self: self._priors_history[-1] if self._priors_history else None,
+ lambda self, v: self._priors_history.append(v),
+ )
+ return orchestrator_cal, orchestrator_prod
+
+ def test_copies_production_condition_growth_and_clears_pinned(self):
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ theta_values = np.array([[0.5]])
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values)
+
+ new_priors = orchestrator_cal._priors
+ assert new_priors.growth.condition_growth.growth_k_hyper_loc == 2.5
+ assert new_priors.growth.condition_growth.k_hyper_loc_scale == 0.4
+ # pinned cleared regardless of production's contents
+ assert new_priors.growth.condition_growth.pinned == {}
+
+ def test_copies_growth_transition_and_clears_pinned(self):
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ theta_values = np.array([[0.5]])
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values)
+
+ new_priors = orchestrator_cal._priors
+ assert new_priors.growth.growth_transition.pre_t_hyper_loc == 3.5
+ assert new_priors.growth.growth_transition.pinned == {}
+
+ def test_handles_growth_transition_without_pinned(self):
+ # The "instant" growth_transition has no pinned field at all.
+ orchestrator_cal, orchestrator_prod = self._make_models(gt_has_pinned=False)
+ theta_values = np.array([[0.5]])
+ # Should not raise.
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values)
+ new_priors = orchestrator_cal._priors
+ assert new_priors.growth.growth_transition.pre_t_hyper_loc == 3.5
+
+ def test_pins_activity_hyperparams_to_prior_locs(self):
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ theta_values = np.array([[0.5]])
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values)
+ pinned = orchestrator_cal._priors.growth.activity.pinned
+ # Both suffixes from _PINNED_COMPONENTS["activity"] populated.
+ for suffix, _ in _PINNED_COMPONENTS["activity"]:
+ assert suffix in pinned
+
+ def test_pins_ln_cfu0_hyperparams_single_class(self):
+ # Single-class case: array fields of length 1 must produce
+ # "hyper_loc_0" and "hyper_scale_0" in the pinned dict, matching
+ # the site suffix "{name}_hyper_loc_0" used by define_model in
+ # the hierarchical ln_cfu0 component.
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, np.array([[0.5]]))
+ pinned = orchestrator_cal._priors.growth.ln_cfu0.pinned
+ assert "hyper_loc_0" in pinned, f"pinned keys: {list(pinned)}"
+ assert "hyper_scale_0" in pinned, f"pinned keys: {list(pinned)}"
+ assert pinned["hyper_loc_0"] == pytest.approx(5.0)
+ assert pinned["hyper_scale_0"] == pytest.approx(1.0)
+
+ def test_pins_ln_cfu0_hyperparams_multi_class(self):
+ # Multi-class case: 2-element arrays must produce per-class entries
+ # "hyper_loc_0", "hyper_loc_1", "hyper_scale_0", "hyper_scale_1".
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ # Replace the ln_cfu0 component in orchestrator_cal.priors (the
+ # MagicMock attribute that _inject_calibration_priors reads).
+ two_class_ln_cfu0 = _FakeLnCfu0(
+ ln_cfu0_hyper_loc_locs=jnp.array([5.0, 7.0]),
+ ln_cfu0_hyper_scale_locs=jnp.array([1.0, 2.0]),
+ )
+ cal_priors = orchestrator_cal.priors
+ new_growth = cal_priors.growth.replace(ln_cfu0=two_class_ln_cfu0)
+ orchestrator_cal.priors = cal_priors.replace(growth=new_growth)
+
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, np.array([[0.5]]))
+ pinned = orchestrator_cal._priors.growth.ln_cfu0.pinned
+ for key, expected in [
+ ("hyper_loc_0", 5.0), ("hyper_loc_1", 7.0),
+ ("hyper_scale_0", 1.0), ("hyper_scale_1", 2.0),
+ ]:
+ assert key in pinned, f"missing '{key}'; pinned keys: {list(pinned)}"
+ assert pinned[key] == pytest.approx(expected)
+
+ def test_dk_geno_not_in_pinned_components(self):
+ # dk_geno uses the fixed component during calibration (all zeros).
+ # Pinning hyperparameters is meaningless for a fixed component, so
+ # dk_geno must not appear in _PINNED_COMPONENTS.
+ assert "dk_geno" not in _PINNED_COMPONENTS
+
+ def test_dk_geno_not_pinned_after_inject(self):
+ # _inject_calibration_priors must leave dk_geno.pinned untouched
+ # (empty) because the fixed component has no hyperparameters.
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, np.array([[0.5]]))
+ assert orchestrator_cal._priors.growth.dk_geno.pinned == {}
+
+ def test_theta_values_are_set(self):
+ orchestrator_cal, orchestrator_prod = self._make_models()
+ # theta_values is now (T, C, G); _inject_calibration_priors is shape-agnostic.
+ theta_values = np.array([[[0.1, 0.9], [0.3, 0.7]]]) # shape (1, 2, 2)
+ _inject_calibration_priors(orchestrator_cal, orchestrator_prod, theta_values)
+ result = orchestrator_cal._priors.theta.theta_values
+ assert np.allclose(np.asarray(result), theta_values)
+
+
+# ---------------------------------------------------------------------------
+# run_prefit_calibration — orchestration with everything mocked
+# ---------------------------------------------------------------------------
+
+class TestRunPrefitCalibrationOrchestration:
+
+ def _write_yaml_and_csvs(self, tmp_path):
+ priors = tmp_path / "p.csv"
+ guesses = tmp_path / "g.csv"
+ priors.write_text("parameter,value\n")
+ guesses.write_text("parameter,value\n")
+ cfg = tmp_path / "cfg.yaml"
+ with open(cfg, "w") as fh:
+ yaml.dump({
+ "priors_file": "p.csv",
+ "guesses_file": "g.csv",
+ "data": {"growth": "growth.csv", "binding": "binding.csv"},
+ }, fh)
+ return str(cfg), str(priors), str(guesses)
+
+ def _patch_pipeline(self, mocker, orchestrator_cal=None, params=None,
+ field_mapping=None):
+ """Stub out every heavy callsite of run_prefit_calibration."""
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.read_configuration",
+ return_value=(MagicMock(), {}),
+ )
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli._intersect_data",
+ return_value=(pd.DataFrame(), pd.DataFrame()),
+ )
+ orchestrator_cal = orchestrator_cal or MagicMock()
+ orchestrator_cal.init_params = {}
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli._build_calibration_model",
+ return_value=orchestrator_cal,
+ )
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli._compute_theta_values",
+ return_value=np.array([[[0.5]]]), # shape (T=1, C=1, G=1)
+ )
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli._inject_calibration_priors",
+ )
+ mock_ri = MagicMock()
+ mock_ri._iterations_per_epoch = 1
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.RunInference",
+ return_value=mock_ri,
+ )
+ mock_run_map = mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli._run_calibration_map",
+ return_value=("svi_state", params or {}, True),
+ )
+ mocker.patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli._identify_field_mapping",
+ return_value=field_mapping or {},
+ )
+ return mock_ri, mock_run_map
+
+ def test_raises_when_seed_and_checkpoint_both_none(self, tmp_path, mocker):
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ with pytest.raises(ValueError, match="seed must be provided"):
+ run_prefit_calibration(config_file=cfg)
+
+ def test_returns_run_calibration_map_result(self, tmp_path, mocker):
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ self._patch_pipeline(mocker, params={"alpha_auto_loc": 1.0})
+ result = run_prefit_calibration(config_file=cfg, seed=42)
+ assert result == ("svi_state", {"alpha_auto_loc": 1.0}, True)
+
+ def test_writes_bak_files_when_updates_present(self, tmp_path, mocker):
+ cfg, priors, guesses = self._write_yaml_and_csvs(tmp_path)
+ # Seed the priors CSV with a loc-field row that the Normal-site
+ # update should target, plus an unrelated row to verify we don't
+ # disturb it.
+ pd.DataFrame({
+ "parameter": [
+ "growth.condition_growth.growth_k_hyper_loc",
+ "growth.activity.hyper_loc_loc",
+ ],
+ "value": [1.0, 0.0],
+ }).to_csv(priors, index=False)
+ pd.DataFrame({
+ "parameter": ["condition_growth_growth_k_hyper"],
+ "value": [1.0],
+ "flat_index": [float("nan")],
+ }).to_csv(guesses, index=False)
+
+ params = {"condition_growth_growth_k_hyper_auto_loc": np.float32(9.0)}
+ self._patch_pipeline(
+ mocker,
+ params=params,
+ field_mapping={
+ "condition_growth_growth_k_hyper": {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": "growth_k_hyper_loc",
+ "scale_field": "k_hyper_loc_scale",
+ "is_array": False,
+ },
+ },
+ )
+ run_prefit_calibration(config_file=cfg, seed=1)
+
+ # Both .bak files written, both live CSVs updated.
+ assert os.path.exists(priors + ".bak")
+ assert os.path.exists(guesses + ".bak")
+ new_priors = pd.read_csv(priors).set_index("parameter")["value"]
+ # MAP → loc-field row; unrelated row preserved.
+ assert new_priors["growth.condition_growth.growth_k_hyper_loc"] == 9.0
+ assert new_priors["growth.activity.hyper_loc_loc"] == 0.0
+ new_guesses = pd.read_csv(guesses)
+ assert new_guesses.iloc[0]["value"] == 9.0
+
+ def test_default_seed_used_when_resuming_from_checkpoint(self, tmp_path,
+ mocker):
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ _, mock_run_map = self._patch_pipeline(mocker)
+ # No seed but a checkpoint_file → should not raise.
+ run_prefit_calibration(config_file=cfg, seed=None,
+ checkpoint_file="/tmp/resume.pkl")
+ # checkpoint_file forwarded to _run_calibration_map
+ kwargs = mock_run_map.call_args.kwargs
+ assert kwargs["checkpoint_file"] == "/tmp/resume.pkl"
+
+ def test_optimizer_kwargs_forwarded(self, tmp_path, mocker):
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ _, mock_run_map = self._patch_pipeline(mocker)
+ run_prefit_calibration(
+ config_file=cfg,
+ seed=1,
+ adam_step_size=2.5e-3,
+ adam_final_step_size=2.5e-7,
+ adam_clip_norm=2.0,
+ elbo_num_particles=5,
+ convergence_tolerance=1e-4,
+ convergence_window=20,
+ patience=4,
+ convergence_check_interval=3,
+ checkpoint_interval=25,
+ max_num_epochs=500,
+ init_param_jitter=0.0,
+ )
+ kwargs = mock_run_map.call_args.kwargs
+ assert kwargs["adam_step_size"] == 2.5e-3
+ assert kwargs["adam_final_step_size"] == 2.5e-7
+ assert kwargs["adam_clip_norm"] == 2.0
+ assert kwargs["elbo_num_particles"] == 5
+ assert kwargs["convergence_tolerance"] == 1e-4
+ assert kwargs["convergence_window"] == 20
+ assert kwargs["patience"] == 4
+ assert kwargs["convergence_check_interval"] == 3
+ assert kwargs["checkpoint_interval"] == 25
+ assert kwargs["max_num_epochs"] == 500
+ assert kwargs["init_param_jitter"] == 0.0
+
+ def test_default_out_prefix_is_prefit(self, tmp_path, mocker):
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ _, mock_run_map = self._patch_pipeline(mocker)
+ run_prefit_calibration(config_file=cfg, seed=1)
+ assert mock_run_map.call_args.kwargs["out_prefix"] == "tfs_prefit"
+
+ def test_custom_out_prefix_is_honored(self, tmp_path, mocker):
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ _, mock_run_map = self._patch_pipeline(mocker)
+ run_prefit_calibration(config_file=cfg, seed=1,
+ out_prefix="my_runA")
+ assert mock_run_map.call_args.kwargs["out_prefix"] == "my_runA"
+
+ def test_default_init_param_jitter_is_zero(self, tmp_path, mocker):
+ """Pre-fit should be deterministic given a seed; default jitter is 0."""
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ _, mock_run_map = self._patch_pipeline(mocker)
+ run_prefit_calibration(config_file=cfg, seed=1)
+ assert mock_run_map.call_args.kwargs["init_param_jitter"] == 0.0
+
+ def test_hessian_called_after_map(self, tmp_path, mocker):
+ """compute_hessian_sigmas must be called exactly once after MAP."""
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ mock_ri, _ = self._patch_pipeline(mocker)
+ run_prefit_calibration(config_file=cfg, seed=1)
+ mock_ri.compute_hessian_sigmas.assert_called_once()
+
+ def test_hessian_chunk_size_forwarded(self, tmp_path, mocker):
+ """The hessian_chunk_size kwarg must be passed to compute_hessian_sigmas."""
+ cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
+ mock_ri, _ = self._patch_pipeline(mocker)
+ run_prefit_calibration(config_file=cfg, seed=1, hessian_chunk_size=32)
+ call_kwargs = mock_ri.compute_hessian_sigmas.call_args.kwargs
+ assert call_kwargs.get("hessian_chunk_size") == 32
+
+
+# ---------------------------------------------------------------------------
+# _build_hessian_scale_updates
+# ---------------------------------------------------------------------------
+
+class TestBuildHessianScaleUpdates:
+ """Tests for the Hessian-with-floor scale update builder."""
+
+ def _make_field_mapping(self, sites):
+ """Build a minimal field_mapping for the given list of (site, loc_field) pairs."""
+ out = {}
+ for site_name, loc_field in sites:
+ suffix = loc_field.replace("_loc", "")
+ out[site_name] = {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": loc_field,
+ "scale_field": f"{suffix}_scale",
+ "is_array": True,
+ }
+ return out
+
+ def _make_hessian(self, site_name, sigmas):
+ return {site_name: {"map": np.zeros_like(sigmas), "sigma": np.array(sigmas)}}
+
+ # --- floor application ---
+
+ def test_k_sigma_above_floor_is_kept(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.01, 0.02])
+ g, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ np.testing.assert_allclose(g["condition_growth_k_scales"], [0.01, 0.02])
+
+ def test_k_sigma_below_floor_is_floored(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.0001, 0.0005])
+ g, _ = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert np.all(g["condition_growth_k_scales"] == pytest.approx(0.002))
+
+ def test_m_sigma_uses_m_floor_not_k_floor(self):
+ fm = self._make_field_mapping([("condition_growth_m", "m_loc")])
+ hr = self._make_hessian("condition_growth_m", [0.0003, 0.0008])
+ g, _ = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ np.testing.assert_allclose(g["condition_growth_m_scales"], [0.001, 0.001])
+
+ def test_floor_applied_elementwise(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.001, 0.005])
+ g, _ = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.003, m_scale_floor=0.001,
+ k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ np.testing.assert_allclose(g["condition_growth_k_scales"], [0.003, 0.005])
+
+ # --- guess output ---
+
+ def test_k_scales_in_guess_updates(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.005, 0.007])
+ g, _ = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert "condition_growth_k_scales" in g
+
+ def test_m_scales_in_guess_updates(self):
+ fm = self._make_field_mapping([("condition_growth_m", "m_loc")])
+ hr = self._make_hessian("condition_growth_m", [0.003, 0.004])
+ g, _ = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert "condition_growth_m_scales" in g
+
+ # --- prior output ---
+
+ def test_k_prior_update_uses_max_across_conditions(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.003, 0.007])
+ _, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert p["growth.condition_growth.k_scale"] == pytest.approx(0.007)
+
+ def test_k_prior_update_respects_floor(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.0001, 0.0002])
+ _, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert p["growth.condition_growth.k_scale"] == pytest.approx(0.002)
+
+ def test_m_prior_goes_to_m_scale_plus_not_m_scale(self):
+ fm = self._make_field_mapping([("condition_growth_m", "m_loc")])
+ hr = self._make_hessian("condition_growth_m", [0.004, 0.006])
+ _, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert "growth.condition_growth.m_scale_plus" in p
+ assert "growth.condition_growth.m_scale" not in p
+
+ def test_m_scale_plus_is_max_floored(self):
+ fm = self._make_field_mapping([("condition_growth_m", "m_loc")])
+ hr = self._make_hessian("condition_growth_m", [0.009, 0.002])
+ _, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert p["growth.condition_growth.m_scale_plus"] == pytest.approx(0.009)
+
+ # --- edge cases ---
+
+ def test_site_missing_from_hessian_skipped(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ g, p = _build_hessian_scale_updates(fm, {}, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert g == {} and p == {}
+
+ def test_non_array_site_skipped(self):
+ fm = {
+ "condition_growth_k_hyper": {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": "k_hyper_loc",
+ "scale_field": "k_hyper_scale",
+ "is_array": False,
+ }
+ }
+ hr = {"condition_growth_k_hyper": {"map": np.array(0.02), "sigma": np.array(0.0001)}}
+ g, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert g == {} and p == {}
+
+ def test_unknown_loc_field_skipped(self):
+ fm = {
+ "condition_growth_tau": {
+ "component": "condition_growth",
+ "dist_class": "Normal",
+ "loc_field": "tau_loc",
+ "scale_field": "tau_scale",
+ "is_array": True,
+ }
+ }
+ hr = {"condition_growth_tau": {"map": np.zeros(2), "sigma": np.ones(2) * 0.1}}
+ g, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert g == {} and p == {}
+
+ def test_joint_k_and_m_both_updated(self):
+ fm = self._make_field_mapping([
+ ("condition_growth_k", "k_loc"),
+ ("condition_growth_m", "m_loc"),
+ ])
+ hr = {
+ "condition_growth_k": {"map": np.zeros(2), "sigma": np.array([0.005, 0.008])},
+ "condition_growth_m": {"map": np.zeros(2), "sigma": np.array([0.004, 0.002])},
+ }
+ g, p = _build_hessian_scale_updates(fm, hr, k_scale_floor=0.002, m_scale_floor=0.001, k_scale_ceiling=0.1, m_scale_ceiling=0.01)
+ assert "condition_growth_k_scales" in g
+ assert "condition_growth_m_scales" in g
+ assert "growth.condition_growth.k_scale" in p
+ assert "growth.condition_growth.m_scale_plus" in p
+
+ # --- ceiling application ---
+
+ def test_k_sigma_above_ceiling_is_capped(self):
+ """A degenerate Hessian giving sigma >> ceiling must not loosen the prior."""
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [23.5, 31.6])
+ g, p = _build_hessian_scale_updates(
+ fm, hr, k_scale_floor=0.002, m_scale_floor=0.001,
+ k_scale_ceiling=0.1, m_scale_ceiling=0.01,
+ )
+ assert np.all(g["condition_growth_k_scales"] == pytest.approx(0.1))
+ assert p["growth.condition_growth.k_scale"] == pytest.approx(0.1)
+
+ def test_m_sigma_above_ceiling_is_capped(self):
+ fm = self._make_field_mapping([("condition_growth_m", "m_loc")])
+ hr = self._make_hessian("condition_growth_m", [15.0, 22.0])
+ g, p = _build_hessian_scale_updates(
+ fm, hr, k_scale_floor=0.002, m_scale_floor=0.001,
+ k_scale_ceiling=0.1, m_scale_ceiling=0.01,
+ )
+ assert np.all(g["condition_growth_m_scales"] == pytest.approx(0.01))
+ assert p["growth.condition_growth.m_scale_plus"] == pytest.approx(0.01)
+
+ def test_sigma_between_floor_and_ceiling_is_unchanged(self):
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.004, 0.006])
+ g, p = _build_hessian_scale_updates(
+ fm, hr, k_scale_floor=0.002, m_scale_floor=0.001,
+ k_scale_ceiling=0.1, m_scale_ceiling=0.01,
+ )
+ np.testing.assert_allclose(g["condition_growth_k_scales"], [0.004, 0.006])
+ assert p["growth.condition_growth.k_scale"] == pytest.approx(0.006)
+
+ def test_ceiling_below_floor_raises_or_clips_to_ceiling(self):
+ """When ceiling < floor (degenerate config), clip still works (clip clamps to ceiling)."""
+ fm = self._make_field_mapping([("condition_growth_k", "k_loc")])
+ hr = self._make_hessian("condition_growth_k", [0.05])
+ g, _ = _build_hessian_scale_updates(
+ fm, hr, k_scale_floor=0.05, m_scale_floor=0.001,
+ k_scale_ceiling=0.05, m_scale_ceiling=0.01,
+ )
+ # floor == ceiling == 0.05 → result must be 0.05
+ assert g["condition_growth_k_scales"][0] == pytest.approx(0.05)
+
+ # --- default constant sanity ---
+
+ def test_default_floors_are_sensible_constants(self):
+ """Floor constants must be positive and k_floor >= m_floor."""
+ assert _DEFAULT_K_SCALE_FLOOR > 0
+ assert _DEFAULT_M_SCALE_FLOOR > 0
+ assert _DEFAULT_K_SCALE_FLOOR >= _DEFAULT_M_SCALE_FLOOR
+
+ def test_default_ceilings_above_floors(self):
+ """Ceilings must be strictly above floors so the valid range is non-empty."""
+ assert _DEFAULT_K_SCALE_CEILING > _DEFAULT_K_SCALE_FLOOR
+ assert _DEFAULT_M_SCALE_CEILING > _DEFAULT_M_SCALE_FLOOR
+
+ def test_default_ceilings_match_linear_defaults(self):
+ """Ceilings should match the linear component's default prior scales."""
+ from tfscreen.tfmodel.generative.components.growth.linear import get_hyperparameters
+ h = get_hyperparameters()
+ assert _DEFAULT_K_SCALE_CEILING == pytest.approx(h["k_scale"])
+ assert _DEFAULT_M_SCALE_CEILING == pytest.approx(h["m_scale_plus"])
+
+
+# ---------------------------------------------------------------------------
+# main() CLI plumbing
+# ---------------------------------------------------------------------------
+
+class TestPrefitMainCLI:
+
+ def test_main_invokes_run_prefit_with_required_args(self, tmp_path):
+ """`main()` should drive the wrapper through generalized_main."""
+ argv = ["cal.yaml", "--seed", "13"]
+ with patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.run_prefit_calibration",
+ autospec=True,
+ ) as mock_run, \
+ patch("sys.argv", ["tfs-prefit-calibration"] + argv):
+ main()
+
+ assert mock_run.call_count == 1
+ kwargs = mock_run.call_args.kwargs
+ assert kwargs["config_file"] == "cal.yaml"
+ assert kwargs["seed"] == 13
+
+ def test_main_forwards_custom_out_prefix(self, tmp_path):
+ argv = [
+ "cal.yaml",
+ "--seed", "0",
+ "--out_prefix", "calibration_runA",
+ ]
+ with patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.run_prefit_calibration",
+ autospec=True,
+ ) as mock_run, \
+ patch("sys.argv", ["tfs-prefit-calibration"] + argv):
+ main()
+ assert mock_run.call_args.kwargs["out_prefix"] == "calibration_runA"
+
+ def test_main_forwards_checkpoint_file(self, tmp_path):
+ argv = [
+ "cal.yaml",
+ "--seed", "0",
+ "--checkpoint_file", "/tmp/ck.pkl",
+ ]
+ with patch(
+ "tfscreen.tfmodel.scripts"
+ ".prefit_calibration_cli.run_prefit_calibration",
+ autospec=True,
+ ) as mock_run, \
+ patch("sys.argv", ["tfs-prefit-calibration"] + argv):
+ main()
+ assert mock_run.call_args.kwargs["checkpoint_file"] == "/tmp/ck.pkl"
diff --git a/tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py b/tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py
deleted file mode 100644
index 09871fde..00000000
--- a/tests/tfscreen/tfmodel/scripts/test_run_prefit_calibration.py
+++ /dev/null
@@ -1,1733 +0,0 @@
-"""
-Tests for run_prefit_calibration.py.
-
-The pre-fit script builds an in-process calibration ModelOrchestrator,
-runs MAP, computes Hessian-based per-site sigmas, and then writes
-in-place updates (with .bak backups) into the production priors and
-guesses CSVs. These tests exercise the helper functions directly and
-mock the heavy machinery (read_configuration, ModelOrchestrator, RunInference)
-when testing the orchestration in run_prefit_calibration / main.
-"""
-import dataclasses
-import os
-import shutil
-import sys
-from unittest.mock import MagicMock, patch
-
-import numpy as np
-import pandas as pd
-import pytest
-import yaml
-import jax
-import jax.numpy as jnp
-
-import flax.struct as fstruct
-
-from tfscreen.tfmodel.scripts.run_prefit_calibration_cli import (
- _apply_guesses_updates,
- _apply_priors_updates,
- _build_calibration_model,
- _build_csv_updates,
- _compute_growth_pred_std,
- _compute_theta_values,
- _csv_row_name,
- _identify_field_mapping,
- _inject_calibration_priors,
- _intersect_data,
- _make_calibration_plots,
- _make_correlation_plot,
- _resolve_csv_paths,
- _write_calibration_stats,
- main,
- run_prefit_calibration,
- _CALIBRATION_OVERRIDES,
- _PINNED_COMPONENTS,
-)
-
-
-# ---------------------------------------------------------------------------
-# Lightweight stand-ins for the flax pytree dataclasses used by ModelPriors.
-#
-# The real component-prior dataclasses are flax.struct.dataclass instances
-# that expose a .replace() method. We define a few minimal substitutes
-# here so the unit tests don't need to instantiate a full ModelOrchestrator.
-# ---------------------------------------------------------------------------
-
-@fstruct.dataclass
-class _FakeCondGrowth:
- growth_k_hyper_loc: float = 0.0
- k_hyper_loc_scale: float = 1.0
- growth_k_hyper_scale_loc: float = 0.5
- growth_m_hyper_loc: float = 0.0
- m_hyper_loc_scale: float = 1.0
- growth_m_hyper_scale_loc: float = 0.5
- pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
-
-
-@fstruct.dataclass
-class _FakeGrowthTransition:
- pre_t_hyper_loc: float = 0.0
- pre_t_hyper_loc_scale: float = 1.0
- pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
-
-
-@fstruct.dataclass
-class _FakeGTNoPinned:
- pre_t_hyper_loc: float = 0.0
- pre_t_hyper_loc_scale: float = 1.0
- # No pinned field; mirrors the "instant" growth_transition variant.
-
-
-@fstruct.dataclass
-class _FakeActivity:
- hyper_loc_loc: float = 0.0
- hyper_scale_loc: float = 1.0
- pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
-
-
-@fstruct.dataclass
-class _FakeDkGeno:
- hyper_loc_loc: float = 0.0
- hyper_scale_loc: float = 1.0
- hyper_shift_loc: float = 0.5
- pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
-
-
-@fstruct.dataclass
-class _FakeLnCfu0:
- ln_cfu0_hyper_loc_loc: float = 5.0
- ln_cfu0_hyper_scale_loc: float = 1.0
- pinned: dict = fstruct.field(default_factory=dict, pytree_node=False)
-
-
-@fstruct.dataclass
-class _FakeTheta:
- theta_values: object = None # array placeholder
-
-
-@fstruct.dataclass
-class _FakeGrowth:
- condition_growth: _FakeCondGrowth
- growth_transition: _FakeGrowthTransition
- activity: _FakeActivity
- dk_geno: _FakeDkGeno
- ln_cfu0: _FakeLnCfu0
-
-
-@fstruct.dataclass
-class _FakePriors:
- growth: _FakeGrowth
- theta: _FakeTheta
-
-
-def _make_fake_priors(gt_has_pinned=True):
- gt = _FakeGrowthTransition() if gt_has_pinned else _FakeGTNoPinned()
- return _FakePriors(
- growth=_FakeGrowth(
- condition_growth=_FakeCondGrowth(),
- growth_transition=gt,
- activity=_FakeActivity(),
- dk_geno=_FakeDkGeno(),
- ln_cfu0=_FakeLnCfu0(),
- ),
- theta=_FakeTheta(),
- )
-
-
-# ---------------------------------------------------------------------------
-# _intersect_data
-# ---------------------------------------------------------------------------
-
-class TestIntersectData:
-
- def _basic_growth(self):
- return pd.DataFrame({
- "genotype": ["wt", "wt", "A1T", "G2P"],
- "titrant_name": ["IPTG"] * 4,
- "titrant_conc": [0.0, 1.0, 0.0, 0.0],
- "ln_cfu": [1.0, 2.0, 3.0, 4.0],
- })
-
- def _basic_binding(self):
- return pd.DataFrame({
- "genotype": ["wt", "wt", "A1T", "M3L"],
- "titrant_name": ["IPTG"] * 4,
- "titrant_conc": [0.0, 1.0, 0.0, 0.0],
- "theta_obs": [0.1, 0.5, 0.2, 0.3],
- "theta_std": [0.05, 0.05, 0.05, 0.05],
- })
-
- def test_intersection_keeps_only_shared_rows(self):
- g, b = _intersect_data(self._basic_growth(), self._basic_binding())
- # G2P is only in growth, M3L is only in binding; both must drop.
- assert set(g["genotype"].unique()) == {"wt", "A1T"}
- assert set(b["genotype"].unique()) == {"wt", "A1T"}
-
- def test_intersection_returns_copies_not_views(self):
- gdf = self._basic_growth()
- bdf = self._basic_binding()
- g, _ = _intersect_data(gdf, bdf)
- # Mutating the returned frame must not alter the original.
- g.loc[g.index[0], "ln_cfu"] = -999.0
- assert gdf.loc[0, "ln_cfu"] != -999.0
-
- def test_intersection_preserves_per_titrant_resolution(self):
- # (wt, IPTG, 1.0) and (wt, IPTG, 0.0) must each be evaluated
- # separately — we cannot just intersect on genotype.
- growth = pd.DataFrame({
- "genotype": ["wt", "wt"],
- "titrant_name": ["IPTG", "IPTG"],
- "titrant_conc": [0.0, 1.0],
- "ln_cfu": [1.0, 2.0],
- })
- binding = pd.DataFrame({
- "genotype": ["wt"],
- "titrant_name": ["IPTG"],
- "titrant_conc": [1.0],
- "theta_obs": [0.5],
- "theta_std": [0.05],
- })
- g, b = _intersect_data(growth, binding)
- assert len(g) == 1
- assert g["titrant_conc"].iloc[0] == 1.0
-
- def test_empty_intersection_raises(self):
- growth = pd.DataFrame({
- "genotype": ["wt"],
- "titrant_name": ["IPTG"],
- "titrant_conc": [0.0],
- "ln_cfu": [1.0],
- })
- binding = pd.DataFrame({
- "genotype": ["A1T"],
- "titrant_name": ["IPTG"],
- "titrant_conc": [0.0],
- "theta_obs": [0.5],
- "theta_std": [0.05],
- })
- with pytest.raises(ValueError, match="intersection is empty"):
- _intersect_data(growth, binding)
-
- def test_missing_column_in_growth_raises(self):
- growth = pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"]})
- binding = self._basic_binding()
- with pytest.raises(ValueError, match="growth_df is missing"):
- _intersect_data(growth, binding)
-
- def test_missing_column_in_binding_raises(self):
- growth = self._basic_growth()
- binding = pd.DataFrame({"genotype": ["wt"], "titrant_name": ["IPTG"]})
- with pytest.raises(ValueError, match="binding_df is missing"):
- _intersect_data(growth, binding)
-
-
-# ---------------------------------------------------------------------------
-# _compute_theta_values
-# ---------------------------------------------------------------------------
-
-class TestComputeThetaValues:
-
- def _make_gm_cal(self, titrant_names, titrant_concs):
- """Return a MagicMock with the binding_tm shape used by the helper."""
- gm = MagicMock()
- # Order matters; the helper reads names then concs by index lookup.
- gm.binding_tm.tensor_dim_names = ["titrant_name", "titrant_conc"]
- gm.binding_tm.tensor_dim_labels = [list(titrant_names),
- list(titrant_concs)]
- return gm
-
- def test_inverse_variance_weighted_mean(self):
- gm = self._make_gm_cal(["IPTG"], [1.0])
- # Two genotypes: theta=0.2 (sigma=0.1) and theta=0.8 (sigma=0.01)
- # Inverse-variance weights heavily favor the second, so the
- # consensus should be ~0.8 (not the unweighted 0.5).
- binding_df = pd.DataFrame({
- "titrant_name": ["IPTG", "IPTG"],
- "titrant_conc": [1.0, 1.0],
- "theta_obs": [0.2, 0.8],
- "theta_std": [0.1, 0.01],
- })
- theta = np.asarray(_compute_theta_values(gm, binding_df))
- assert theta.shape == (1, 1)
- # Closer to 0.8 than to the midpoint 0.5
- assert theta[0, 0] > 0.7
-
- def test_falls_back_to_plain_mean_when_all_stds_zero(self):
- gm = self._make_gm_cal(["IPTG"], [1.0])
- binding_df = pd.DataFrame({
- "titrant_name": ["IPTG", "IPTG"],
- "titrant_conc": [1.0, 1.0],
- "theta_obs": [0.3, 0.7],
- "theta_std": [0.0, 0.0], # invalid weights → fallback to mean
- })
- theta = np.asarray(_compute_theta_values(gm, binding_df))
- assert theta[0, 0] == pytest.approx(0.5)
-
- def test_unobserved_cell_defaults_to_midpoint(self):
- gm = self._make_gm_cal(["IPTG"], [0.0, 1.0])
- # Only the IPTG=1.0 cell is observed.
- binding_df = pd.DataFrame({
- "titrant_name": ["IPTG"],
- "titrant_conc": [1.0],
- "theta_obs": [0.9],
- "theta_std": [0.01],
- })
- theta = np.asarray(_compute_theta_values(gm, binding_df))
- assert theta[0, 0] == pytest.approx(0.5)
- assert theta[0, 1] > 0.85
-
- def test_clip_to_open_interval(self):
- gm = self._make_gm_cal(["IPTG"], [0.0, 1.0])
- binding_df = pd.DataFrame({
- "titrant_name": ["IPTG", "IPTG"],
- "titrant_conc": [0.0, 1.0],
- "theta_obs": [0.0, 1.0],
- "theta_std": [0.01, 0.01],
- })
- theta = np.asarray(_compute_theta_values(gm, binding_df))
- # Zero and one would blow up the downstream logit; expect clipping.
- assert theta[0, 0] > 0.0
- assert theta[0, 1] < 1.0
-
-
-
-
-# ---------------------------------------------------------------------------
-# _csv_row_name / _build_csv_updates
-# ---------------------------------------------------------------------------
-
-class TestCsvRowName:
-
- def test_basic_row_name(self):
- assert _csv_row_name("condition_growth", "growth_k_hyper_loc") == \
- "growth.condition_growth.growth_k_hyper_loc"
-
-
-class TestBuildCsvUpdates:
-
- def test_normal_site_writes_loc_and_scale(self):
- field_mapping = {
- "growth_k_hyper": {
- "component": "condition_growth",
- "dist_class": "Normal",
- "loc_field": "growth_k_hyper_loc",
- "scale_field": "k_hyper_loc_scale",
- "is_array": False,
- },
- }
- hessian_results = {
- "growth_k_hyper": {"map": np.float32(2.5),
- "sigma": np.float32(0.7)},
- }
- prior_updates, guess_updates = _build_csv_updates(field_mapping,
- hessian_results)
- # MAP → loc_field; sigma → scale_field
- assert prior_updates[
- "growth.condition_growth.growth_k_hyper_loc"] == pytest.approx(2.5)
- assert prior_updates[
- "growth.condition_growth.k_hyper_loc_scale"] == pytest.approx(0.7)
- assert guess_updates["growth_k_hyper"] == pytest.approx(2.5)
-
- def test_halfnormal_site_recenters_on_map(self):
- field_mapping = {
- "k_hyper_scale_loc": {
- "component": "condition_growth",
- "dist_class": "HalfNormal",
- "scale_field": "growth_k_hyper_scale_loc",
- "is_array": False,
- },
- }
- hessian_results = {
- "k_hyper_scale_loc": {"map": np.float32(0.42),
- "sigma": np.float32(0.09)},
- }
- prior_updates, guess_updates = _build_csv_updates(field_mapping,
- hessian_results)
- # HalfNormal has only one parameter (its scale). We recenter the
- # prior on the MAP point — sigma is dropped because there is no
- # second parameter to consume it.
- assert prior_updates == {
- "growth.condition_growth.growth_k_hyper_scale_loc": pytest.approx(0.42),
- }
- assert guess_updates["k_hyper_scale_loc"] == pytest.approx(0.42)
-
- def test_skips_sites_without_hessian_result(self):
- field_mapping = {
- "missing_site": {
- "component": "condition_growth",
- "dist_class": "Normal",
- "loc_field": "growth_k_hyper_loc",
- "scale_field": "k_hyper_loc_scale",
- "is_array": False,
- },
- }
- prior_updates, guess_updates = _build_csv_updates(field_mapping, {})
- assert prior_updates == {}
- assert guess_updates == {}
-
- def test_array_site_writes_locs_to_guesses(self):
- """Simple-prior array sites write per-condition MAP values to guesses."""
- field_mapping = {
- "condition_growth_k": {
- "component": "condition_growth",
- "dist_class": "Normal",
- "loc_field": "k_loc",
- "scale_field": "k_scale",
- "is_array": True,
- },
- }
- hessian_results = {
- "condition_growth_k": {
- "map": np.array([1.0, 2.0, 3.0]),
- "sigma": np.array([0.1, 0.2, 0.3]),
- },
- }
- prior_updates, guess_updates = _build_csv_updates(field_mapping,
- hessian_results)
- # Array sites do NOT update priors
- assert prior_updates == {}
- # Guess locs get the per-condition MAP array
- assert "condition_growth_k_locs" in guess_updates
- assert np.allclose(guess_updates["condition_growth_k_locs"],
- [1.0, 2.0, 3.0])
-
-
-# ---------------------------------------------------------------------------
-# _apply_priors_updates / _apply_guesses_updates
-# ---------------------------------------------------------------------------
-
-class TestApplyPriorsUpdates:
-
- def _write_csv(self, tmp_path, rows):
- path = tmp_path / "priors.csv"
- pd.DataFrame(rows).to_csv(path, index=False)
- return str(path)
-
- def test_writes_bak_and_updates_only_matching_rows(self, tmp_path):
- path = self._write_csv(tmp_path, [
- {"parameter": "growth.condition_growth.growth_k_hyper_loc",
- "value": 1.0},
- {"parameter": "growth.activity.hyper_loc_loc",
- "value": 0.0},
- ])
- updates = {"growth.condition_growth.growth_k_hyper_loc": 9.0}
- _apply_priors_updates(path, updates)
-
- bak = pd.read_csv(path + ".bak")
- new = pd.read_csv(path)
- # Backup retains the original
- assert bak.set_index("parameter")["value"][
- "growth.condition_growth.growth_k_hyper_loc"] == 1.0
- # Live file has the updated value
- assert new.set_index("parameter")["value"][
- "growth.condition_growth.growth_k_hyper_loc"] == 9.0
- # Untouched rows are preserved verbatim
- assert new.set_index("parameter")["value"][
- "growth.activity.hyper_loc_loc"] == 0.0
-
- def test_no_updates_is_noop(self, tmp_path):
- path = self._write_csv(tmp_path, [
- {"parameter": "x", "value": 1.0},
- ])
- _apply_priors_updates(path, {})
- # No .bak, no rewrite.
- assert not os.path.exists(path + ".bak")
-
- def test_warns_on_missing_row(self, tmp_path, capsys):
- path = self._write_csv(tmp_path, [
- {"parameter": "x", "value": 1.0},
- ])
- _apply_priors_updates(path, {"y_does_not_exist": 5.0})
- err = capsys.readouterr().err
- assert "no matching row" in err
- assert "y_does_not_exist" in err
-
- def test_raises_on_malformed_csv(self, tmp_path):
- bad = tmp_path / "bad.csv"
- pd.DataFrame({"foo": [1, 2]}).to_csv(bad, index=False)
- with pytest.raises(ValueError, match="missing required"):
- _apply_priors_updates(str(bad), {"x": 1.0})
-
-
-class TestApplyGuessesUpdates:
-
- def _write_csv(self, tmp_path, rows):
- path = tmp_path / "guesses.csv"
- pd.DataFrame(rows).to_csv(path, index=False)
- return str(path)
-
- def test_writes_bak_and_updates_only_scalar_rows(self, tmp_path):
- # Two rows for the same parameter: one scalar (flat_index NaN),
- # one array entry (flat_index=0). Only the scalar must be
- # overwritten.
- path = self._write_csv(tmp_path, [
- {"parameter": "growth_k_hyper", "value": 1.0,
- "flat_index": float("nan")},
- {"parameter": "growth_k_offset", "value": 0.0, "flat_index": 0.0},
- {"parameter": "growth_k_offset", "value": 0.0, "flat_index": 1.0},
- ])
- updates = {"growth_k_hyper": 7.7,
- "growth_k_offset": 999.0}
- _apply_guesses_updates(path, updates)
-
- new = pd.read_csv(path)
- # Scalar row updated
- scalar_row = new[(new["parameter"] == "growth_k_hyper")
- & (new["flat_index"].isna())]
- assert scalar_row["value"].iloc[0] == pytest.approx(7.7)
- # Array rows untouched (the update key matched no scalar row)
- offset_rows = new[new["parameter"] == "growth_k_offset"]
- assert (offset_rows["value"] == 0.0).all()
-
- # .bak written
- assert os.path.exists(path + ".bak")
-
- def test_no_updates_is_noop(self, tmp_path):
- path = self._write_csv(tmp_path, [
- {"parameter": "x", "value": 1.0, "flat_index": float("nan")},
- ])
- _apply_guesses_updates(path, {})
- assert not os.path.exists(path + ".bak")
-
- def test_updates_array_rows_by_flat_index(self, tmp_path):
- path = self._write_csv(tmp_path, [
- {"parameter": "condition_growth_k_locs", "value": 0.0, "flat_index": 0.0},
- {"parameter": "condition_growth_k_locs", "value": 0.0, "flat_index": 1.0},
- {"parameter": "condition_growth_k_locs", "value": 0.0, "flat_index": 2.0},
- {"parameter": "other_param", "value": 9.9, "flat_index": float("nan")},
- ])
- _apply_guesses_updates(path, {"condition_growth_k_locs": np.array([1.1, 2.2, 3.3])})
- new = pd.read_csv(path)
- k_rows = new[new["parameter"] == "condition_growth_k_locs"].sort_values("flat_index")
- assert k_rows["value"].tolist() == pytest.approx([1.1, 2.2, 3.3])
- # Unrelated row preserved
- assert new[new["parameter"] == "other_param"]["value"].iloc[0] == pytest.approx(9.9)
- assert os.path.exists(path + ".bak")
-
- def test_treats_all_rows_scalar_when_no_flat_index_column(self, tmp_path):
- path = self._write_csv(tmp_path, [
- {"parameter": "alpha", "value": 1.0},
- ])
- _apply_guesses_updates(path, {"alpha": 4.5})
- new = pd.read_csv(path)
- assert new.set_index("parameter")["value"]["alpha"] == 4.5
-
- def test_raises_on_malformed_csv(self, tmp_path):
- bad = tmp_path / "bad.csv"
- pd.DataFrame({"foo": [1]}).to_csv(bad, index=False)
- with pytest.raises(ValueError, match="missing required"):
- _apply_guesses_updates(str(bad), {"x": 1.0})
-
-
-# ---------------------------------------------------------------------------
-# _resolve_csv_paths
-# ---------------------------------------------------------------------------
-
-class TestResolveCsvPaths:
-
- def _write_yaml(self, path, body):
- with open(path, "w") as fh:
- yaml.dump(body, fh)
-
- def test_returns_resolved_paths(self, tmp_path):
- priors = tmp_path / "p.csv"
- guesses = tmp_path / "g.csv"
- priors.write_text("parameter,value\n")
- guesses.write_text("parameter,value\n")
- cfg = tmp_path / "cfg.yaml"
- self._write_yaml(cfg, {"priors_file": "p.csv",
- "guesses_file": "g.csv"})
- p, g = _resolve_csv_paths(str(cfg))
- assert os.path.abspath(p) == str(priors)
- assert os.path.abspath(g) == str(guesses)
-
- def test_absolute_paths_are_honored(self, tmp_path):
- priors = tmp_path / "abs_p.csv"
- guesses = tmp_path / "abs_g.csv"
- priors.write_text("parameter,value\n")
- guesses.write_text("parameter,value\n")
- cfg = tmp_path / "cfg.yaml"
- self._write_yaml(cfg, {"priors_file": str(priors),
- "guesses_file": str(guesses)})
- p, g = _resolve_csv_paths(str(cfg))
- assert p == str(priors)
- assert g == str(guesses)
-
- def test_missing_priors_key_raises(self, tmp_path):
- cfg = tmp_path / "cfg.yaml"
- self._write_yaml(cfg, {"guesses_file": "g.csv"})
- with pytest.raises(ValueError, match="priors_file"):
- _resolve_csv_paths(str(cfg))
-
- def test_missing_guesses_key_raises(self, tmp_path):
- cfg = tmp_path / "cfg.yaml"
- self._write_yaml(cfg, {"priors_file": "p.csv"})
- with pytest.raises(ValueError, match="guesses_file"):
- _resolve_csv_paths(str(cfg))
-
- def test_missing_priors_file_raises(self, tmp_path):
- cfg = tmp_path / "cfg.yaml"
- guesses = tmp_path / "g.csv"
- guesses.write_text("parameter,value\n")
- self._write_yaml(cfg, {"priors_file": "p_missing.csv",
- "guesses_file": "g.csv"})
- with pytest.raises(FileNotFoundError, match="Priors file not found"):
- _resolve_csv_paths(str(cfg))
-
- def test_missing_guesses_file_raises(self, tmp_path):
- cfg = tmp_path / "cfg.yaml"
- priors = tmp_path / "p.csv"
- priors.write_text("parameter,value\n")
- self._write_yaml(cfg, {"priors_file": "p.csv",
- "guesses_file": "g_missing.csv"})
- with pytest.raises(FileNotFoundError, match="Guesses file not found"):
- _resolve_csv_paths(str(cfg))
-
-
-# ---------------------------------------------------------------------------
-# _build_calibration_model — verify it forwards the right overrides
-# ---------------------------------------------------------------------------
-
-class TestBuildCalibrationModel:
-
- def test_overrides_replace_production_components(self):
- gm_prod = MagicMock()
- gm_prod.settings = {
- "theta": "categorical_geno", # → simple
- "activity": "horseshoe_geno", # → hierarchical
- "dk_geno": "fixed", # → hierarchical
- "ln_cfu0": "fixed", # → hierarchical
- "transformation": "logit_norm", # → single
- "theta_growth_noise": "beta", # → zero
- "theta_binding_noise": "beta", # → zero
- "condition_growth": "linear", # passthrough
- "growth_transition": "instant", # passthrough
- "batch_size": 7,
- "spiked_genotypes": ["WT"],
- }
-
- with patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.ModelOrchestrator"
- ) as MockGM:
- MockGM.return_value = MagicMock()
- growth_df = pd.DataFrame()
- binding_df = pd.DataFrame()
- _build_calibration_model(gm_prod, growth_df, binding_df)
-
- kwargs = MockGM.call_args.kwargs
- # Overrides applied
- for k, v in _CALIBRATION_OVERRIDES.items():
- assert kwargs[k] == v
- # Spiked genotypes dropped (calibration only sees the intersection)
- assert kwargs["spiked_genotypes"] is None
- # batch_size pulled out of settings and passed positionally
- assert kwargs["batch_size"] == 7
- # condition_growth and growth_transition flow through unchanged
- assert kwargs["condition_growth"] == "linear"
- assert kwargs["growth_transition"] == "instant"
- # binding_weight must be 1.0 regardless of production value so the
- # calibration MAP learns the binding→growth linkage without the
- # production upweighting drowning the binding signal.
- assert kwargs["binding_weight"] == 1.0
-
- def test_binding_weight_reset_from_large_production_value(self):
- # Reproduces the weighting bug: production YAML stores a large
- # binding_weight (N_growth / N_binding); the calibration model
- # must use 1.0 instead.
- gm_prod = MagicMock()
- gm_prod.settings = {
- "theta": "categorical_geno",
- "activity": "horseshoe_geno",
- "dk_geno": "fixed",
- "ln_cfu0": "fixed",
- "transformation": "logit_norm",
- "theta_growth_noise": "beta",
- "theta_binding_noise": "beta",
- "condition_growth": "linear",
- "growth_transition": "instant",
- "batch_size": None,
- "spiked_genotypes": None,
- "binding_weight": 250.0, # typical large production value
- }
-
- with patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.ModelOrchestrator"
- ) as MockGM:
- MockGM.return_value = MagicMock()
- _build_calibration_model(gm_prod, pd.DataFrame(), pd.DataFrame())
-
- assert MockGM.call_args.kwargs["binding_weight"] == 1.0
-
- def test_does_not_mutate_production_settings(self):
- # Ensure we work on a copy.
- gm_prod = MagicMock()
- gm_prod.settings = {"theta": "categorical_geno",
- "activity": "horseshoe_geno",
- "dk_geno": "fixed",
- "ln_cfu0": "fixed",
- "transformation": "logit_norm",
- "theta_growth_noise": "beta",
- "theta_binding_noise": "beta",
- "spiked_genotypes": None,
- "batch_size": None,
- "binding_weight": 150.0}
- original_settings = dict(gm_prod.settings)
-
- with patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.ModelOrchestrator"
- ) as MockGM:
- MockGM.return_value = MagicMock()
- _build_calibration_model(gm_prod, pd.DataFrame(), pd.DataFrame())
-
- assert gm_prod.settings == original_settings
-
-
-# ---------------------------------------------------------------------------
-# _inject_calibration_priors — end-to-end pin / copy behaviour
-# ---------------------------------------------------------------------------
-
-class TestInjectCalibrationPriors:
-
- def _make_models(self, gt_has_pinned=True):
- gm_cal = MagicMock()
- gm_prod = MagicMock()
- gm_cal._priors_history = []
-
- cal_priors = _make_fake_priors(gt_has_pinned=gt_has_pinned)
- prod_priors = _make_fake_priors(gt_has_pinned=True)
- # Tweak production scalars so we can verify they get copied.
- prod_priors = prod_priors.replace(
- growth=prod_priors.growth.replace(
- condition_growth=prod_priors.growth.condition_growth.replace(
- growth_k_hyper_loc=2.5,
- k_hyper_loc_scale=0.4,
- ),
- growth_transition=prod_priors.growth.growth_transition.replace(
- pre_t_hyper_loc=3.5,
- ),
- ),
- )
- gm_cal.priors = cal_priors
- gm_prod.priors = prod_priors
-
- # Capture all writes to _priors so the test can introspect them.
- def _setter(value):
- gm_cal._priors_history.append(value)
- type(gm_cal)._priors = property(
- lambda self: self._priors_history[-1] if self._priors_history else None,
- lambda self, v: self._priors_history.append(v),
- )
- return gm_cal, gm_prod
-
- def test_copies_production_condition_growth_and_clears_pinned(self):
- gm_cal, gm_prod = self._make_models()
- theta_values = np.array([[0.5]])
- _inject_calibration_priors(gm_cal, gm_prod, theta_values)
-
- new_priors = gm_cal._priors
- assert new_priors.growth.condition_growth.growth_k_hyper_loc == 2.5
- assert new_priors.growth.condition_growth.k_hyper_loc_scale == 0.4
- # pinned cleared regardless of production's contents
- assert new_priors.growth.condition_growth.pinned == {}
-
- def test_copies_growth_transition_and_clears_pinned(self):
- gm_cal, gm_prod = self._make_models()
- theta_values = np.array([[0.5]])
- _inject_calibration_priors(gm_cal, gm_prod, theta_values)
-
- new_priors = gm_cal._priors
- assert new_priors.growth.growth_transition.pre_t_hyper_loc == 3.5
- assert new_priors.growth.growth_transition.pinned == {}
-
- def test_handles_growth_transition_without_pinned(self):
- # The "instant" growth_transition has no pinned field at all.
- gm_cal, gm_prod = self._make_models(gt_has_pinned=False)
- theta_values = np.array([[0.5]])
- # Should not raise.
- _inject_calibration_priors(gm_cal, gm_prod, theta_values)
- new_priors = gm_cal._priors
- assert new_priors.growth.growth_transition.pre_t_hyper_loc == 3.5
-
- def test_pins_activity_hyperparams_to_prior_locs(self):
- gm_cal, gm_prod = self._make_models()
- theta_values = np.array([[0.5]])
- _inject_calibration_priors(gm_cal, gm_prod, theta_values)
- pinned = gm_cal._priors.growth.activity.pinned
- # Both suffixes from _PINNED_COMPONENTS["activity"] populated.
- for suffix, _ in _PINNED_COMPONENTS["activity"]:
- assert suffix in pinned
-
- def test_pins_dk_geno_and_ln_cfu0(self):
- gm_cal, gm_prod = self._make_models()
- _inject_calibration_priors(gm_cal, gm_prod, np.array([[0.5]]))
- for comp in ("dk_geno", "ln_cfu0"):
- pinned = getattr(gm_cal._priors.growth, comp).pinned
- for suffix, _ in _PINNED_COMPONENTS[comp]:
- assert suffix in pinned
-
- def test_theta_values_are_set(self):
- gm_cal, gm_prod = self._make_models()
- theta_values = np.array([[0.1, 0.9]])
- _inject_calibration_priors(gm_cal, gm_prod, theta_values)
- result = gm_cal._priors.theta.theta_values
- assert np.allclose(np.asarray(result), theta_values)
-
-
-# ---------------------------------------------------------------------------
-# run_prefit_calibration — orchestration with everything mocked
-# ---------------------------------------------------------------------------
-
-class TestRunPrefitCalibrationOrchestration:
-
- def _write_yaml_and_csvs(self, tmp_path):
- priors = tmp_path / "p.csv"
- guesses = tmp_path / "g.csv"
- priors.write_text("parameter,value\n")
- guesses.write_text("parameter,value\n")
- cfg = tmp_path / "cfg.yaml"
- with open(cfg, "w") as fh:
- yaml.dump({
- "priors_file": "p.csv",
- "guesses_file": "g.csv",
- "data": {"growth": "growth.csv", "binding": "binding.csv"},
- }, fh)
- return str(cfg), str(priors), str(guesses)
-
- def _patch_pipeline(self, mocker, gm_cal=None, params=None,
- hessian_results=None, field_mapping=None):
- """Stub out every heavy callsite of run_prefit_calibration."""
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.read_configuration",
- return_value=(MagicMock(), {}),
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._intersect_data",
- return_value=(pd.DataFrame(), pd.DataFrame()),
- )
- gm_cal = gm_cal or MagicMock()
- gm_cal.init_params = {}
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._build_calibration_model",
- return_value=gm_cal,
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._compute_theta_values",
- return_value=np.array([[0.5]]),
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._inject_calibration_priors",
- )
- mock_ri = MagicMock()
- mock_ri._iterations_per_epoch = 1
- mock_ri.compute_hessian_sigmas.return_value = hessian_results or {}
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.RunInference",
- return_value=mock_ri,
- )
- mock_run_map = mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._run_calibration_map",
- return_value=("svi_state", params or {}, True),
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._identify_field_mapping",
- return_value=field_mapping or {},
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._compute_growth_pred_std",
- return_value=None,
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._make_calibration_plots",
- )
- return mock_ri, mock_run_map
-
- def test_raises_when_seed_and_checkpoint_both_none(self, tmp_path, mocker):
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- with pytest.raises(ValueError, match="seed must be provided"):
- run_prefit_calibration(config_file=cfg)
-
- def test_returns_run_calibration_map_result(self, tmp_path, mocker):
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- self._patch_pipeline(mocker, params={"alpha_auto_loc": 1.0})
- result = run_prefit_calibration(config_file=cfg, seed=42)
- assert result == ("svi_state", {"alpha_auto_loc": 1.0}, True)
-
- def test_writes_bak_files_when_updates_present(self, tmp_path, mocker):
- cfg, priors, guesses = self._write_yaml_and_csvs(tmp_path)
- # Seed the priors CSV with both the loc-field and scale-field rows
- # that the Normal-site update should target, plus an unrelated
- # row to verify we don't disturb it.
- pd.DataFrame({
- "parameter": [
- "growth.condition_growth.growth_k_hyper_loc",
- "growth.condition_growth.k_hyper_loc_scale",
- "growth.activity.hyper_loc_loc",
- ],
- "value": [1.0, 1.0, 0.0],
- }).to_csv(priors, index=False)
- pd.DataFrame({
- "parameter": ["growth_k_hyper"],
- "value": [1.0],
- "flat_index": [float("nan")],
- }).to_csv(guesses, index=False)
-
- self._patch_pipeline(
- mocker,
- params={"growth_k_hyper_auto_loc": 9.0},
- hessian_results={
- "growth_k_hyper": {"map": np.float32(9.0),
- "sigma": np.float32(0.5)},
- },
- field_mapping={
- "growth_k_hyper": {
- "component": "condition_growth",
- "dist_class": "Normal",
- "loc_field": "growth_k_hyper_loc",
- "scale_field": "k_hyper_loc_scale",
- "is_array": False,
- },
- },
- )
- run_prefit_calibration(config_file=cfg, seed=1)
-
- # Both .bak files written, both live CSVs updated.
- assert os.path.exists(priors + ".bak")
- assert os.path.exists(guesses + ".bak")
- new_priors = pd.read_csv(priors).set_index("parameter")["value"]
- # MAP → loc-field row; Hessian sigma → scale-field row;
- # unrelated row preserved.
- assert new_priors["growth.condition_growth.growth_k_hyper_loc"] == 9.0
- assert new_priors[
- "growth.condition_growth.k_hyper_loc_scale"] == pytest.approx(0.5)
- assert new_priors["growth.activity.hyper_loc_loc"] == 0.0
- new_guesses = pd.read_csv(guesses)
- assert new_guesses.iloc[0]["value"] == 9.0
-
- def test_default_seed_used_when_resuming_from_checkpoint(self, tmp_path,
- mocker):
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- _, mock_run_map = self._patch_pipeline(mocker)
- # No seed but a checkpoint_file → should not raise.
- run_prefit_calibration(config_file=cfg, seed=None,
- checkpoint_file="/tmp/resume.pkl")
- # checkpoint_file forwarded to _run_calibration_map
- kwargs = mock_run_map.call_args.kwargs
- assert kwargs["checkpoint_file"] == "/tmp/resume.pkl"
-
- def test_optimizer_kwargs_forwarded(self, tmp_path, mocker):
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- _, mock_run_map = self._patch_pipeline(mocker)
- run_prefit_calibration(
- config_file=cfg,
- seed=1,
- adam_step_size=2.5e-3,
- adam_final_step_size=2.5e-7,
- adam_clip_norm=2.0,
- elbo_num_particles=5,
- convergence_tolerance=1e-4,
- convergence_window=20,
- patience=4,
- convergence_check_interval=3,
- checkpoint_interval=25,
- max_num_epochs=500,
- init_param_jitter=0.0,
- )
- kwargs = mock_run_map.call_args.kwargs
- assert kwargs["adam_step_size"] == 2.5e-3
- assert kwargs["adam_final_step_size"] == 2.5e-7
- assert kwargs["adam_clip_norm"] == 2.0
- assert kwargs["elbo_num_particles"] == 5
- assert kwargs["convergence_tolerance"] == 1e-4
- assert kwargs["convergence_window"] == 20
- assert kwargs["patience"] == 4
- assert kwargs["convergence_check_interval"] == 3
- assert kwargs["checkpoint_interval"] == 25
- assert kwargs["max_num_epochs"] == 500
- assert kwargs["init_param_jitter"] == 0.0
-
- def test_default_out_prefix_is_prefit(self, tmp_path, mocker):
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- _, mock_run_map = self._patch_pipeline(mocker)
- run_prefit_calibration(config_file=cfg, seed=1)
- assert mock_run_map.call_args.kwargs["out_prefix"] == "tfs_prefit"
-
- def test_custom_out_prefix_is_honored(self, tmp_path, mocker):
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- _, mock_run_map = self._patch_pipeline(mocker)
- run_prefit_calibration(config_file=cfg, seed=1,
- out_prefix="my_runA")
- assert mock_run_map.call_args.kwargs["out_prefix"] == "my_runA"
-
- def test_default_init_param_jitter_is_zero(self, tmp_path, mocker):
- """Pre-fit should be deterministic given a seed; default jitter is 0."""
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- _, mock_run_map = self._patch_pipeline(mocker)
- run_prefit_calibration(config_file=cfg, seed=1)
- assert mock_run_map.call_args.kwargs["init_param_jitter"] == 0.0
-
- def test_compute_growth_pred_std_called_with_ri_and_params(self, tmp_path,
- mocker):
- """_compute_growth_pred_std is called with the RunInference instance
- and the exact params dict returned by _run_calibration_map."""
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- mock_params = {"alpha_auto_loc": np.float32(2.0)}
- mock_ri, _ = self._patch_pipeline(mocker, params=mock_params)
-
- # Override the stub so we can inspect calls.
- mock_std_fn = mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._compute_growth_pred_std",
- return_value=None,
- )
- run_prefit_calibration(config_file=cfg, seed=1)
-
- assert mock_std_fn.call_count == 1
- call_args = mock_std_fn.call_args
- assert call_args.args[0] is mock_ri # first positional: ri
- assert call_args.args[1] is mock_params # second positional: params
-
- def test_growth_pred_std_forwarded_to_make_calibration_plots(self, tmp_path,
- mocker):
- """The growth_pred_std returned by _compute_growth_pred_std is passed
- as the growth_pred_std keyword argument to _make_calibration_plots."""
- cfg, _, _ = self._write_yaml_and_csvs(tmp_path)
- sentinel = np.ones((1, 2, 1, 1, 1, 1, 2))
- self._patch_pipeline(mocker)
-
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._compute_growth_pred_std",
- return_value=sentinel,
- )
- mock_plots = mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._make_calibration_plots",
- )
- run_prefit_calibration(config_file=cfg, seed=1)
-
- plots_kwargs = mock_plots.call_args.kwargs
- assert plots_kwargs.get("growth_pred_std") is sentinel
-
-
-# ---------------------------------------------------------------------------
-# main() CLI plumbing
-# ---------------------------------------------------------------------------
-
-class TestPrefitMainCLI:
-
- def test_main_invokes_run_prefit_with_required_args(self, tmp_path):
- """`main()` should drive the wrapper through generalized_main."""
- argv = ["cal.yaml", "--seed", "13"]
- with patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.run_prefit_calibration",
- autospec=True,
- ) as mock_run, \
- patch("sys.argv", ["tfs-prefit-calibration"] + argv):
- main()
-
- assert mock_run.call_count == 1
- kwargs = mock_run.call_args.kwargs
- assert kwargs["config_file"] == "cal.yaml"
- assert kwargs["seed"] == 13
-
- def test_main_forwards_custom_out_prefix(self, tmp_path):
- argv = [
- "cal.yaml",
- "--seed", "0",
- "--out_prefix", "calibration_runA",
- ]
- with patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.run_prefit_calibration",
- autospec=True,
- ) as mock_run, \
- patch("sys.argv", ["tfs-prefit-calibration"] + argv):
- main()
- assert mock_run.call_args.kwargs["out_prefix"] == "calibration_runA"
-
- def test_main_forwards_checkpoint_file(self, tmp_path):
- argv = [
- "cal.yaml",
- "--seed", "0",
- "--checkpoint_file", "/tmp/ck.pkl",
- ]
- with patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.run_prefit_calibration",
- autospec=True,
- ) as mock_run, \
- patch("sys.argv", ["tfs-prefit-calibration"] + argv):
- main()
- assert mock_run.call_args.kwargs["checkpoint_file"] == "/tmp/ck.pkl"
-
-
-# ---------------------------------------------------------------------------
-# _make_calibration_plots
-# ---------------------------------------------------------------------------
-
-def _build_fake_gm_cal(n_rep=1, n_t=3, n_cp=1, n_cs=1, n_tn=1, n_tc=1,
- n_geno=2, include_ln_cfu0=True):
- """
- Build a minimal MagicMock stand-in for ModelOrchestrator that satisfies
- every attribute access in _make_calibration_plots.
-
- Tensor shape convention: (R, T, CP, CS, TN, TC, G).
- """
- shape = (n_rep, n_t, n_cp, n_cs, n_tn, n_tc, n_geno)
-
- # growth data tensors
- good_mask = np.ones(shape, dtype=bool)
- good_mask[0, 0, 0, 0, 0, 0, 0] = False # one padded cell to exercise masking
- t_pre_tensor = np.full(shape, 5.0)
- t_sel_tensor = np.zeros(shape)
- for t_i in range(n_t):
- t_sel_tensor[:, t_i, ...] = float(t_i + 1)
- ln_cfu_obs = np.ones(shape) * 10.0
- ln_cfu_std = np.ones(shape) * 0.5
-
- gd = MagicMock()
- gd.good_mask = good_mask
- gd.t_pre = t_pre_tensor
- gd.t_sel = t_sel_tensor
- gd.ln_cfu = ln_cfu_obs
- gd.ln_cfu_std = ln_cfu_std
-
- data = MagicMock()
- data.growth = gd
- data.num_genotype = n_geno
-
- # TensorManager mock
- dim_names = ["replicate", "condition_pre", "condition_sel",
- "titrant_name", "titrant_conc", "genotype"]
- dim_labels = [
- [f"rep{i}" for i in range(n_rep)],
- [f"cp{i}" for i in range(n_cp)],
- [f"cs{i}" for i in range(n_cs)],
- [f"tn{i}" for i in range(n_tn)],
- [float(i) for i in range(n_tc)],
- [f"geno{i}" for i in range(n_geno)],
- ]
-
- # Build a realistic tm.df with all index columns needed for the merge in
- # _write_predictions_csv, and condition_sel for cs_name_map building.
- rows = []
- for r_i in range(n_rep):
- for t_i in range(n_t):
- for cp_i in range(n_cp):
- for cs_i in range(n_cs):
- for tn_i in range(n_tn):
- for tc_i in range(n_tc):
- for g_i in range(n_geno):
- rows.append({
- "replicate_idx": r_i,
- "condition_pre_idx": cp_i,
- "condition_sel_idx": cs_i,
- "titrant_name_idx": tn_i,
- "titrant_conc_idx": tc_i,
- "genotype_idx": g_i,
- "t_sel": float(t_i + 1),
- "condition_sel": f"cs{cs_i}",
- })
- df = pd.DataFrame(rows)
-
- tm = MagicMock()
- tm.tensor_dim_names = dim_names
- tm.tensor_dim_labels = dim_labels
- tm.df = df
-
- # MAP samples returned by Predictive
- growth_pred_val = np.ones((1,) + shape) * 10.0 # (samples, R,T,CP,CS,TN,TC,G)
- map_samples = {"growth_pred": growth_pred_val}
- if include_ln_cfu0:
- ln_cfu0_shape = (1, n_rep, n_cp, n_geno) # (samples, R, CP, G)
- map_samples["ln_cfu0"] = np.ones(ln_cfu0_shape) * 9.0
-
- # Configure the full_data mock returned by get_batch. The production code
- # accesses full_data.growth.{map_condition_pre, map_condition_sel, t_pre}
- # to build the fine-grid data; they must be proper numpy arrays so that
- # np.asarray() and slicing work correctly inside _bc_t.
- map_cond = np.zeros(shape, dtype=int) # constant-zero indices are fine
- full_data_gd = MagicMock()
- full_data_gd.map_condition_pre = map_cond
- full_data_gd.map_condition_sel = map_cond
- full_data_gd.t_pre = t_pre_tensor
- full_data_gd.good_mask = good_mask
- # .replace() is called but its return value is only passed to pred_fn,
- # which is fully mocked; so returning a plain MagicMock is fine.
- full_data_mock = MagicMock()
- full_data_mock.growth = full_data_gd
-
- gm_cal = MagicMock()
- gm_cal.data = data
- gm_cal.priors = MagicMock()
- gm_cal.growth_tm = tm
- gm_cal.jax_model = MagicMock()
- gm_cal.get_batch = MagicMock(return_value=full_data_mock)
-
- return gm_cal, map_samples
-
-
-class TestMakeCalibrationPlots:
- """Tests for _make_calibration_plots."""
-
- def _patch_predictive(self, mocker, map_samples):
- """Patch numpyro.infer.Predictive and AutoDelta inside the function.
-
- _make_calibration_plots calls pred_fn twice:
- 1st call — checked for 'growth_pred' presence and provides ln_cfu0.
- 2nd call — fine-grid predictions; growth_pred must have T_FINE=50
- time points so that np.concatenate([..., t_fine_1d],
- [..., y_smooth_sel]) does not mis-match.
- """
- T_FINE = 50
- if "growth_pred" in map_samples:
- orig_gp = map_samples["growth_pred"] # (1, R, n_t, CP, CS, TN, TC, G)
- shape_fine = list(orig_gp.shape)
- shape_fine[2] = T_FINE
- map_samples_fine = {**map_samples,
- "growth_pred": np.ones(shape_fine) * 10.0}
- else:
- map_samples_fine = map_samples
-
- mock_pred_instance = MagicMock(
- side_effect=[map_samples, map_samples_fine]
- )
- mock_pred_cls = mocker.patch(
- "numpyro.infer.Predictive",
- return_value=mock_pred_instance,
- )
- mocker.patch("numpyro.infer.autoguide.AutoDelta")
- mocker.patch("jax.random.PRNGKey", return_value=0)
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.jnp.arange",
- return_value=np.arange(2),
- )
- return mock_pred_cls
-
- def test_creates_one_pdf_per_genotype(self, tmp_path, mocker):
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=3)
- self._patch_predictive(mocker, map_samples)
-
- _make_calibration_plots(gm_cal, params={}, out_prefix=str(tmp_path / "run"))
-
- pdfs = sorted(tmp_path.glob("*.pdf"))
- assert len(pdfs) == 3
-
- def test_pdf_names_use_out_prefix_and_genotype(self, tmp_path, mocker):
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=2)
- self._patch_predictive(mocker, map_samples)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "myrun"))
-
- pdf_names = {p.name for p in tmp_path.glob("*.pdf")}
- assert "myrun_calib_geno0.pdf" in pdf_names
- assert "myrun_calib_geno1.pdf" in pdf_names
-
- def test_skips_genotype_with_no_valid_observations(self, tmp_path, mocker):
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=2)
-
- # Make all observations invalid for genotype 0 only.
- good_mask = np.asarray(gm_cal.data.growth.good_mask)
- good_mask[..., 0] = False
- gm_cal.data.growth.good_mask = good_mask
-
- self._patch_predictive(mocker, map_samples)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- pdfs = {p.name for p in tmp_path.glob("*.pdf")}
- assert "run_calib_geno0.pdf" not in pdfs
- assert "run_calib_geno1.pdf" in pdfs
-
- def test_warns_and_returns_when_growth_pred_missing(self, tmp_path, mocker,
- capsys):
- gm_cal, _ = _build_fake_gm_cal()
- # Exclude growth_pred so the function should bail out early.
- map_samples_no_pred = {"ln_cfu0": np.ones((1, 1, 1, 2)) * 9.0}
- self._patch_predictive(mocker, map_samples_no_pred)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- assert not list(tmp_path.glob("*.pdf"))
- captured = capsys.readouterr()
- assert "growth_pred" in captured.err
-
- def test_handles_missing_ln_cfu0_gracefully(self, tmp_path, mocker):
- """When ln_cfu0 is absent from map_samples, falls back to val_at_t0."""
- gm_cal, map_samples = _build_fake_gm_cal(include_ln_cfu0=False)
- self._patch_predictive(mocker, map_samples)
-
- # Should not raise even without ln_cfu0.
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- assert len(list(tmp_path.glob("*.pdf"))) == 2
-
- def test_single_time_point_does_not_crash(self, tmp_path, mocker):
- """polyfit requires ≥2 pts; single-point case uses fallback branch."""
- gm_cal, map_samples = _build_fake_gm_cal(n_t=1)
- # Remove the masked cell so both genotypes have a valid observation.
- good_mask = np.asarray(gm_cal.data.growth.good_mask)
- good_mask[:] = True
- gm_cal.data.growth.good_mask = good_mask
- self._patch_predictive(mocker, map_samples)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- # A PDF must still be produced for each genotype.
- assert len(list(tmp_path.glob("*.pdf"))) == 2
-
- def test_matplotlib_unavailable_warns_and_returns(self, tmp_path, mocker,
- capsys):
- """If matplotlib cannot be imported, a warning is emitted and no PDF
- is written."""
- gm_cal, map_samples = _build_fake_gm_cal()
- self._patch_predictive(mocker, map_samples)
-
- # Simulate ImportError for matplotlib by making the import raise.
- real_import = __builtins__["__import__"] if isinstance(__builtins__, dict) \
- else __import__
-
- def fake_import(name, *args, **kwargs):
- if name == "matplotlib":
- raise ImportError("no matplotlib")
- return real_import(name, *args, **kwargs)
-
- mocker.patch("builtins.__import__", side_effect=fake_import)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- assert not list(tmp_path.glob("*.pdf"))
- captured = capsys.readouterr()
- assert "matplotlib" in captured.err
-
- def test_writes_predictions_csv(self, tmp_path, mocker):
- """A CSV named {out_prefix}_calib_growth_df.csv is written with a
- ln_cfu_pred column populated for every valid observation."""
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=2, n_t=3)
- self._patch_predictive(mocker, map_samples)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- csv_path = tmp_path / "run_calib_growth_df.csv"
- assert csv_path.exists(), "predictions CSV was not created"
-
- result = pd.read_csv(csv_path)
- assert "ln_cfu_pred" in result.columns
-
- # The default fixture sets growth_pred to 10.0 everywhere; the masked
- # cell (r=0, t=0, cp=0, cs=0, tn=0, tc=0, g=0) stays NaN because it
- # is excluded from good_mask and therefore from pred_lookup.
- assert result["ln_cfu_pred"].notna().any()
- valid_preds = result["ln_cfu_pred"].dropna()
- assert np.allclose(valid_preds.values, 10.0)
-
- def test_writes_ln_cfu_pred_std_when_provided(self, tmp_path, mocker):
- """ln_cfu_pred_std column appears in the CSV when growth_pred_std is given."""
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=2, n_t=3)
- self._patch_predictive(mocker, map_samples)
-
- good_mask = np.asarray(gm_cal.data.growth.good_mask)
- std_array = np.full(good_mask.shape, 0.25)
-
- _make_calibration_plots(
- gm_cal, params={},
- out_prefix=str(tmp_path / "run"),
- growth_pred_std=std_array,
- )
-
- result = pd.read_csv(tmp_path / "run_calib_growth_df.csv")
- assert "ln_cfu_pred_std" in result.columns
- valid_std = result.dropna(subset=["ln_cfu_pred_std"])["ln_cfu_pred_std"]
- assert np.allclose(valid_std.values, 0.25)
-
- def test_no_ln_cfu_pred_std_column_when_not_provided(self, tmp_path, mocker):
- """When growth_pred_std is None the CSV must not contain the std column."""
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=2, n_t=3)
- self._patch_predictive(mocker, map_samples)
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- result = pd.read_csv(tmp_path / "run_calib_growth_df.csv")
- assert "ln_cfu_pred_std" not in result.columns
-
-
-# ---------------------------------------------------------------------------
-# _compute_growth_pred_std
-# ---------------------------------------------------------------------------
-
-class TestComputeGrowthPredStd:
- """Tests for _compute_growth_pred_std.
-
- The function wraps heavy JAX machinery (Hessian, Cholesky, vmap, Predictive).
- Tests that require the Hessian/Predictive path stub those out via
- ``_patch_heavy_machinery``; the empty-params case is exercised directly.
- """
-
- # ------------------------------------------------------------------
- # Helpers
- # ------------------------------------------------------------------
-
- def _patch_heavy_machinery(self, mocker, n_params, n_samples, pred_output):
- """Stub JAX/NumPyro operations so the function reaches Predictive.
-
- * ``jax.device_put`` → identity passthrough
- * ``jax.hessian`` → returns the n_params × n_params identity
- * ``jax.random.normal`` → zeros, so all Laplace samples equal MAP
- * ``trace`` / ``seed`` (in the module) → empty model trace, meaning
- no constrained transforms and no per-site bijections applied
- * ``numpyro.infer.Predictive`` → returns ``pred_output``
- """
- mocker.patch("jax.device_put", side_effect=lambda x: x)
- hess = jnp.eye(n_params)
- # jax.hessian(fn) returns a callable; that callable(flat_map) returns
- # the matrix. Use a MagicMock so the two-call chain works correctly.
- mock_hess_fn = MagicMock(return_value=hess)
- mocker.patch("jax.hessian", return_value=mock_hess_fn)
- mocker.patch(
- "jax.random.normal",
- return_value=jnp.zeros((n_samples, n_params)),
- )
- mock_traced = MagicMock()
- mock_traced.get_trace.return_value = {}
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.trace",
- return_value=mock_traced,
- )
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.seed",
- return_value=MagicMock(),
- )
- mock_pred_inst = MagicMock(return_value=pred_output)
- mocker.patch("numpyro.infer.Predictive", return_value=mock_pred_inst)
-
- def _make_inputs(self, n_params=2):
- """Build (ri, gm_cal, params) mocks for _compute_growth_pred_std."""
- ri = MagicMock()
- ri.get_key.return_value = jax.random.PRNGKey(0)
- gm_cal = MagicMock()
- gm_cal.data.num_genotype = 3
- gm_cal.get_batch.return_value = MagicMock()
- gm_cal.priors = MagicMock()
- params = {f"site{i}_auto_loc": np.float32(float(i))
- for i in range(n_params)}
- return ri, gm_cal, params
-
- # ------------------------------------------------------------------
- # Tests
- # ------------------------------------------------------------------
-
- def test_returns_none_when_params_empty(self):
- """No _auto_loc params → early-exit None without touching JAX."""
- ri = MagicMock()
- gm_cal = MagicMock()
- gm_cal.data.num_genotype = 2
- gm_cal.get_batch.return_value = MagicMock()
- assert _compute_growth_pred_std(ri, {}, gm_cal) is None
-
- def test_returns_none_when_growth_pred_absent(self, mocker):
- """When Predictive returns no growth_pred site, result must be None."""
- n_params, n_samples = 2, 4
- ri, gm_cal, params = self._make_inputs(n_params)
- self._patch_heavy_machinery(
- mocker, n_params, n_samples,
- pred_output={"other_site": np.ones((n_samples, 3))},
- )
- assert _compute_growth_pred_std(ri, params, gm_cal,
- n_samples=n_samples) is None
-
- def test_returns_array_matching_growth_pred_shape(self, mocker):
- """Output shape must equal the (R,T,CP,CS,TN,TC,G) tensor shape."""
- n_params, n_samples = 2, 4
- tensor_shape = (1, 3, 1, 1, 1, 1, 2)
- ri, gm_cal, params = self._make_inputs(n_params)
- growth_pred = np.ones((n_samples,) + tensor_shape) * 5.0
- self._patch_heavy_machinery(
- mocker, n_params, n_samples,
- pred_output={"growth_pred": growth_pred},
- )
- result = _compute_growth_pred_std(ri, params, gm_cal,
- n_samples=n_samples)
- assert result is not None
- assert result.shape == tensor_shape
-
- def test_std_is_zero_for_identical_samples(self, mocker):
- """All Laplace samples identical → elementwise std must be 0."""
- n_params, n_samples = 2, 6
- tensor_shape = (1, 2, 1, 1, 1, 1, 3)
- ri, gm_cal, params = self._make_inputs(n_params)
- growth_pred = np.ones((n_samples,) + tensor_shape) * 7.0
- self._patch_heavy_machinery(
- mocker, n_params, n_samples,
- pred_output={"growth_pred": growth_pred},
- )
- result = _compute_growth_pred_std(ri, params, gm_cal,
- n_samples=n_samples)
- assert np.allclose(result, 0.0)
-
- def test_std_matches_numpy_std_of_samples(self, mocker):
- """Output must equal np.std(growth_pred_samples, axis=0) exactly."""
- n_params, n_samples = 2, 10
- tensor_shape = (1, 2, 1, 1, 1, 1, 3)
- ri, gm_cal, params = self._make_inputs(n_params)
- rng = np.random.default_rng(42)
- growth_pred = rng.normal(size=(n_samples,) + tensor_shape).astype(np.float32)
- self._patch_heavy_machinery(
- mocker, n_params, n_samples,
- pred_output={"growth_pred": growth_pred},
- )
- result = _compute_growth_pred_std(ri, params, gm_cal,
- n_samples=n_samples)
- # Compare in float64 (the function casts before computing std).
- expected = growth_pred.astype(np.float64).std(axis=0)
- assert np.allclose(result, expected, atol=1e-5)
-
- def test_no_overflow_for_large_growth_pred_values(self, mocker):
- """float32 growth_pred values large enough to overflow when squared
- must produce finite std (not inf/nan) after float64 cast + clipping."""
- n_params, n_samples = 2, 4
- tensor_shape = (1, 2, 1, 1, 1, 1, 2)
- ri, gm_cal, params = self._make_inputs(n_params)
- # Values near the float32 overflow boundary (~1e38); squaring these
- # in float32 would overflow to inf inside np.std.
- growth_pred = np.full((n_samples,) + tensor_shape, 1e20, dtype=np.float32)
- growth_pred[0] = -1e20 # introduce variance so std != 0
- self._patch_heavy_machinery(
- mocker, n_params, n_samples,
- pred_output={"growth_pred": growth_pred},
- )
- result = _compute_growth_pred_std(ri, params, gm_cal,
- n_samples=n_samples)
- assert result is not None
- assert np.all(np.isfinite(result))
-
-
-# ---------------------------------------------------------------------------
-# _write_calibration_stats
-# ---------------------------------------------------------------------------
-
-def _make_stats_df(n=20, include_std=True, seed=7):
- """Return a minimal DataFrame suitable for _write_calibration_stats."""
- rng = np.random.default_rng(seed)
- df = pd.DataFrame({
- "genotype": [f"g{i % 3}" for i in range(n)],
- "ln_cfu": rng.normal(10.0, 1.0, n),
- "ln_cfu_pred": rng.normal(10.0, 1.0, n),
- })
- if include_std:
- df["ln_cfu_pred_std"] = rng.uniform(0.1, 0.5, n)
- return df
-
-
-class TestWriteCalibrationStats:
-
- def test_writes_json_with_expected_keys(self, tmp_path):
- df = _make_stats_df()
- _write_calibration_stats(df, str(tmp_path / "run"))
- json_path = tmp_path / "run_calib_stats.json"
- assert json_path.exists()
- import json
- stats = json.loads(json_path.read_text())
- expected_keys = {
- "pct_success", "rmse", "normalized_rmse", "pearson_r",
- "spearman_r", "r_squared", "mean_error", "coverage_prob",
- "residual_corr", "residual_corr_p_value", "bp_p_value",
- "n_params", "n_obs",
- }
- assert expected_keys == set(stats.keys())
-
- def test_all_values_are_python_float_int_or_null(self, tmp_path):
- """Values must be Python float, int, or null (no numpy scalars or NaN)."""
- df = _make_stats_df()
- _write_calibration_stats(df, str(tmp_path / "run"))
- import json
- stats = json.loads((tmp_path / "run_calib_stats.json").read_text())
- int_keys = {"n_params", "n_obs"}
- for k, v in stats.items():
- if k in int_keys:
- assert v is None or isinstance(v, int), \
- f"Key '{k}' has non-int value {v!r}"
- else:
- assert v is None or isinstance(v, float), \
- f"Key '{k}' has non-float value {v!r}"
- if isinstance(v, float):
- assert np.isfinite(v), f"Key '{k}' is not finite: {v}"
-
- def test_n_params_and_n_obs_written_when_provided(self, tmp_path):
- """n_params and n_obs appear as integers in the JSON when supplied."""
- import json
- df = _make_stats_df(n=10)
- _write_calibration_stats(df, str(tmp_path / "run"), n_params=7, n_obs=10)
- stats = json.loads((tmp_path / "run_calib_stats.json").read_text())
- assert stats["n_params"] == 7
- assert stats["n_obs"] == 10
- assert isinstance(stats["n_params"], int)
- assert isinstance(stats["n_obs"], int)
-
- def test_n_params_and_n_obs_null_when_not_provided(self, tmp_path):
- """n_params and n_obs are null in the JSON when not supplied."""
- import json
- df = _make_stats_df()
- _write_calibration_stats(df, str(tmp_path / "run"))
- stats = json.loads((tmp_path / "run_calib_stats.json").read_text())
- assert stats["n_params"] is None
- assert stats["n_obs"] is None
-
- def test_returns_silently_when_std_column_missing(self, tmp_path):
- """No JSON written when ln_cfu_pred_std is absent."""
- df = _make_stats_df(include_std=False)
- _write_calibration_stats(df, str(tmp_path / "run"))
- assert not (tmp_path / "run_calib_stats.json").exists()
-
- def test_returns_silently_when_pred_column_missing(self, tmp_path):
- """No JSON written when ln_cfu_pred is absent."""
- df = _make_stats_df()
- df = df.drop(columns=["ln_cfu_pred"])
- _write_calibration_stats(df, str(tmp_path / "run"))
- assert not (tmp_path / "run_calib_stats.json").exists()
-
- def test_returns_silently_when_no_valid_rows(self, tmp_path):
- """No JSON written when all rows have NaN predictions."""
- df = _make_stats_df()
- df["ln_cfu_pred"] = np.nan
- _write_calibration_stats(df, str(tmp_path / "run"))
- assert not (tmp_path / "run_calib_stats.json").exists()
-
- def test_nan_stats_serialised_as_null(self, tmp_path):
- """Stats entries that are NaN (e.g. constant inputs) must become null."""
- # Constant ln_cfu makes pearson_r undefined → NaN.
- df = _make_stats_df()
- df["ln_cfu"] = 10.0
- _write_calibration_stats(df, str(tmp_path / "run"))
- import json
- stats = json.loads((tmp_path / "run_calib_stats.json").read_text())
- assert stats["pearson_r"] is None
-
-
-# ---------------------------------------------------------------------------
-# _make_correlation_plot
-# ---------------------------------------------------------------------------
-
-def _make_corr_df(n_genotypes=3, n_per=8, seed=11):
- """Return a minimal DataFrame for _make_correlation_plot."""
- rng = np.random.default_rng(seed)
- rows = []
- for i in range(n_genotypes):
- obs = rng.normal(10.0 + i, 1.0, n_per)
- pred = obs + rng.normal(0, 0.3, n_per)
- for o, p in zip(obs, pred):
- rows.append({"genotype": f"g{i}", "ln_cfu": o, "ln_cfu_pred": p})
- return pd.DataFrame(rows)
-
-
-class TestMakeCorrelationPlot:
-
- def test_writes_pdf(self, tmp_path):
- df = _make_corr_df()
- _make_correlation_plot(df, str(tmp_path / "run"))
- assert (tmp_path / "run_calib_correlation.pdf").exists()
-
- def test_returns_silently_when_genotype_column_missing(self, tmp_path):
- df = _make_corr_df().drop(columns=["genotype"])
- _make_correlation_plot(df, str(tmp_path / "run"))
- assert not (tmp_path / "run_calib_correlation.pdf").exists()
-
- def test_returns_silently_when_no_valid_rows(self, tmp_path):
- df = _make_corr_df()
- df["ln_cfu_pred"] = np.nan
- _make_correlation_plot(df, str(tmp_path / "run"))
- assert not (tmp_path / "run_calib_correlation.pdf").exists()
-
- def test_colors_cycle_for_more_genotypes_than_prop_cycle(self, tmp_path,
- mocker):
- """When n_genotypes > len(prop_cycle) the modulo wraps colours; no
- KeyError or IndexError must be raised."""
- # Default prop_cycle has 10 colours; use 15 genotypes.
- df = _make_corr_df(n_genotypes=15)
- # Should complete without error and produce a PDF.
- _make_correlation_plot(df, str(tmp_path / "run"))
- assert (tmp_path / "run_calib_correlation.pdf").exists()
-
- def test_matplotlib_unavailable_returns_silently(self, tmp_path, mocker):
- df = _make_corr_df()
- real_import = __builtins__["__import__"] if isinstance(__builtins__, dict) \
- else __import__
-
- def fake_import(name, *args, **kwargs):
- if name == "matplotlib.pyplot":
- raise ImportError("no matplotlib")
- return real_import(name, *args, **kwargs)
-
- mocker.patch("builtins.__import__", side_effect=fake_import)
- _make_correlation_plot(df, str(tmp_path / "run"))
- assert not (tmp_path / "run_calib_correlation.pdf").exists()
-
- def test_make_calibration_plots_calls_both_helpers(self, tmp_path, mocker):
- """_make_calibration_plots must call _write_calibration_stats and
- _make_correlation_plot after writing the CSV."""
- mock_stats = mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._write_calibration_stats",
- )
- mock_corr = mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli._make_correlation_plot",
- )
- gm_cal, map_samples = _build_fake_gm_cal(n_geno=2, n_t=2)
-
- # Minimal _patch_predictive inline (avoids importing the fixture helper).
- T_FINE = 50
- orig = map_samples["growth_pred"]
- shape_fine = list(orig.shape); shape_fine[2] = T_FINE
- fine = {**map_samples, "growth_pred": np.ones(shape_fine) * 10.0}
- mock_pred = MagicMock(side_effect=[map_samples, fine])
- mocker.patch("numpyro.infer.Predictive", return_value=mock_pred)
- mocker.patch("numpyro.infer.autoguide.AutoDelta")
- mocker.patch("jax.random.PRNGKey", return_value=0)
- mocker.patch(
- "tfscreen.tfmodel.scripts"
- ".run_prefit_calibration_cli.jnp.arange",
- return_value=np.arange(2),
- )
-
- _make_calibration_plots(gm_cal, params={},
- out_prefix=str(tmp_path / "run"))
-
- assert mock_stats.call_count == 1
- assert mock_corr.call_count == 1
- # Both receive the same merged DataFrame as first argument.
- df_arg_stats = mock_stats.call_args.args[0]
- df_arg_corr = mock_corr.call_args.args[0]
- assert isinstance(df_arg_stats, pd.DataFrame)
- assert isinstance(df_arg_corr, pd.DataFrame)
- assert df_arg_stats is df_arg_corr
- # n_params and n_obs must be forwarded as keyword arguments.
- stats_kwargs = mock_stats.call_args.kwargs
- assert "n_params" in stats_kwargs
- assert "n_obs" in stats_kwargs
- assert isinstance(stats_kwargs["n_params"], int)
- assert isinstance(stats_kwargs["n_obs"], int)
diff --git a/tests/tfscreen/tfmodel/scripts/test_sample_posterior_cli.py b/tests/tfscreen/tfmodel/scripts/test_sample_posterior_cli.py
index 1c3af7be..fe653793 100644
--- a/tests/tfscreen/tfmodel/scripts/test_sample_posterior_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_sample_posterior_cli.py
@@ -180,22 +180,23 @@ def test_map_detection_requires_auto_loc_key(self, tmp_path):
ckpt_path = str(tmp_path / "svi.pkl")
open(ckpt_path, "w").close()
+ ri.restore_svi_from_checkpoint.return_value = (MagicMock(), MagicMock())
+
with patch("tfscreen.tfmodel.scripts.sample_posterior_cli.read_configuration",
return_value=(MagicMock(), {})), \
patch("tfscreen.tfmodel.scripts.sample_posterior_cli.RunInference",
return_value=ri), \
- patch("tfscreen.tfmodel.scripts.sample_posterior_cli.dill") as mock_dill, \
- patch("tfscreen.tfmodel.scripts.sample_posterior_cli._run_svi") as mock_svi:
+ patch("tfscreen.tfmodel.scripts.sample_posterior_cli.dill") as mock_dill:
mock_dill.load.return_value = {"svi_state": MagicMock()}
- mock_svi.side_effect = lambda *a, **kw: open(h5_src, "w").close()
+ ri.get_posteriors.side_effect = lambda **kw: open(h5_src, "w").close()
from tfscreen.tfmodel.scripts.sample_posterior_cli import sample_posterior
sample_posterior("cfg.yaml", ckpt_path,
out_prefix=str(tmp_path / "out"))
ri.get_laplace_posteriors.assert_not_called()
- mock_svi.assert_called_once()
+ ri.get_posteriors.assert_called_once()
# ---------------------------------------------------------------------------
@@ -206,11 +207,9 @@ class TestSamplePosteriorSvi:
def _run_svi_test(self, tmp_path):
ri = MagicMock()
- fake_optim = MagicMock()
- fake_optim.get_params.return_value = {"some_param_mean": np.array(1.0)}
- fake_svi = MagicMock()
- fake_svi.optim = fake_optim
- ri.setup_svi.return_value = fake_svi
+ fake_svi_obj = MagicMock()
+ fake_svi_state = MagicMock()
+ ri.restore_svi_from_checkpoint.return_value = (fake_svi_obj, fake_svi_state)
h5_src = str(tmp_path / "out_tmp_posterior_posterior.h5")
ckpt_path = str(tmp_path / "svi.pkl")
@@ -218,19 +217,18 @@ def _run_svi_test(self, tmp_path):
captured = {}
- def fake_run_svi(ri_obj, init_params, checkpoint_file, out_prefix,
- max_num_epochs, always_get_posterior=False, **kwargs):
- captured["max_num_epochs"] = max_num_epochs
- captured["always_get_posterior"] = always_get_posterior
+ def fake_get_posteriors(**kwargs):
+ captured["get_posteriors_called"] = True
+ captured["get_posteriors_kwargs"] = kwargs
open(h5_src, "w").close()
+ ri.get_posteriors.side_effect = fake_get_posteriors
+
with patch("tfscreen.tfmodel.scripts.sample_posterior_cli.read_configuration",
return_value=(MagicMock(), {})), \
patch("tfscreen.tfmodel.scripts.sample_posterior_cli.RunInference",
return_value=ri), \
- patch("tfscreen.tfmodel.scripts.sample_posterior_cli.dill") as mock_dill, \
- patch("tfscreen.tfmodel.scripts.sample_posterior_cli._run_svi",
- side_effect=fake_run_svi):
+ patch("tfscreen.tfmodel.scripts.sample_posterior_cli.dill") as mock_dill:
mock_dill.load.return_value = {"svi_state": MagicMock()}
@@ -238,15 +236,15 @@ def fake_run_svi(ri_obj, init_params, checkpoint_file, out_prefix,
sample_posterior("cfg.yaml", ckpt_path,
out_prefix=str(tmp_path / "out"))
- return captured
+ return ri, captured
- def test_svi_checkpoint_calls_run_svi_with_zero_epochs(self, tmp_path):
- captured = self._run_svi_test(tmp_path)
- assert captured["max_num_epochs"] == 0
+ def test_svi_calls_restore_svi_from_checkpoint(self, tmp_path):
+ ri, _ = self._run_svi_test(tmp_path)
+ ri.restore_svi_from_checkpoint.assert_called_once()
- def test_svi_always_get_posterior_is_true(self, tmp_path):
- captured = self._run_svi_test(tmp_path)
- assert captured["always_get_posterior"] is True
+ def test_svi_calls_get_posteriors(self, tmp_path):
+ _, captured = self._run_svi_test(tmp_path)
+ assert captured.get("get_posteriors_called") is True
def test_svi_output_renamed(self, tmp_path):
self._run_svi_test(tmp_path)
diff --git a/tests/tfscreen/tfmodel/scripts/test_prior_predictive_cli.py b/tests/tfscreen/tfmodel/scripts/test_sample_prior_cli.py
similarity index 68%
rename from tests/tfscreen/tfmodel/scripts/test_prior_predictive_cli.py
rename to tests/tfscreen/tfmodel/scripts/test_sample_prior_cli.py
index 4718acb5..fd60c1fc 100644
--- a/tests/tfscreen/tfmodel/scripts/test_prior_predictive_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_sample_prior_cli.py
@@ -1,4 +1,4 @@
-"""Tests for prior_predictive_cli.py — draws and writes synthetic prior datasets."""
+"""Tests for sample_prior_cli.py — draws and writes synthetic prior datasets."""
import os
import pytest
import numpy as np
@@ -14,13 +14,13 @@
class TestWriteH5:
def test_creates_file(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import _write_h5
+ from tfscreen.tfmodel.scripts.sample_prior_cli import _write_h5
path = str(tmp_path / "out.h5")
_write_h5(path, {"x": np.ones((5, 4))}, num_draws=5)
assert os.path.isfile(path)
def test_datasets_present(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import _write_h5
+ from tfscreen.tfmodel.scripts.sample_prior_cli import _write_h5
path = str(tmp_path / "out.h5")
data = {"alpha": np.zeros((3, 2)), "beta": np.ones(10)}
_write_h5(path, data, num_draws=3)
@@ -29,7 +29,7 @@ def test_datasets_present(self, tmp_path):
assert "beta" in hf
def test_data_values_correct(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import _write_h5
+ from tfscreen.tfmodel.scripts.sample_prior_cli import _write_h5
path = str(tmp_path / "out.h5")
arr = np.arange(12, dtype=float).reshape(3, 4)
_write_h5(path, {"arr": arr}, num_draws=3)
@@ -37,21 +37,21 @@ def test_data_values_correct(self, tmp_path):
np.testing.assert_array_equal(hf["arr"][:], arr)
def test_num_samples_attr(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import _write_h5
+ from tfscreen.tfmodel.scripts.sample_prior_cli import _write_h5
path = str(tmp_path / "out.h5")
_write_h5(path, {"x": np.ones(5)}, num_draws=7)
with h5py.File(path, "r") as hf:
assert hf.attrs["num_samples"] == 7
def test_gzip_compression(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import _write_h5
+ from tfscreen.tfmodel.scripts.sample_prior_cli import _write_h5
path = str(tmp_path / "out.h5")
_write_h5(path, {"x": np.ones((100, 50))}, num_draws=100)
with h5py.File(path, "r") as hf:
assert hf["x"].compression == "gzip"
def test_1d_array_chunked(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import _write_h5
+ from tfscreen.tfmodel.scripts.sample_prior_cli import _write_h5
path = str(tmp_path / "out.h5")
arr = np.zeros(20)
_write_h5(path, {"flat": arr}, num_draws=20)
@@ -71,17 +71,17 @@ def _mock_draw_prior(self):
return predictions, latent_params
def test_writes_csv_and_h5_for_one_dataset(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
fake_gm = MagicMock()
fake_growth_df = pd.DataFrame({"ln_cfu": [1.0, 2.0]})
latent_params = {"activity": np.array([[1.0]])}
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(fake_gm, {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
return_value=({}, latent_params)), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
return_value=fake_growth_df):
out_prefix = str(tmp_path / "prior")
@@ -91,17 +91,17 @@ def test_writes_csv_and_h5_for_one_dataset(self, tmp_path):
assert os.path.isfile(f"{out_prefix}_000_ground_truth.h5")
def test_writes_multiple_datasets(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
fake_gm = MagicMock()
latent_params = {"activity": np.array([[1.0]])}
fake_growth_df = pd.DataFrame({"ln_cfu": [1.0]})
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(fake_gm, {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
return_value=({}, latent_params)), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
return_value=fake_growth_df):
out_prefix = str(tmp_path / "prior")
@@ -113,16 +113,16 @@ def test_writes_multiple_datasets(self, tmp_path):
def test_zero_padded_index_width(self, tmp_path):
"""Indices must be zero-padded to at least 3 digits."""
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
latent_params = {"activity": np.array([[1.0]])}
fake_growth_df = pd.DataFrame({"ln_cfu": [1.0]})
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(MagicMock(), {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
return_value=({}, latent_params)), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
return_value=fake_growth_df):
out_prefix = str(tmp_path / "prior")
@@ -132,21 +132,21 @@ def test_zero_padded_index_width(self, tmp_path):
def test_each_dataset_uses_different_seed(self, tmp_path):
"""draw_prior must be called with seed+i for dataset i."""
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
seeds_used = []
latent_params = {"activity": np.array([[1.0]])}
fake_growth_df = pd.DataFrame({"ln_cfu": [1.0]})
- def fake_draw_prior(gm, rng_key, num_draws):
+ def fake_draw_prior(orchestrator, rng_key, num_draws):
seeds_used.append(rng_key)
return {}, latent_params
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(MagicMock(), {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
side_effect=fake_draw_prior), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
return_value=fake_growth_df):
out_prefix = str(tmp_path / "prior")
@@ -156,20 +156,20 @@ def fake_draw_prior(gm, rng_key, num_draws):
def test_no_noise_passes_none_rng(self, tmp_path):
"""noise=False must pass noise_rng=None to growth_df_from_prior."""
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
captured = {}
latent_params = {"activity": np.array([[1.0]])}
- def fake_growth_df(gm, params, draw_idx, noise_rng):
+ def fake_growth_df(orchestrator, params, draw_idx, noise_rng):
captured["noise_rng"] = noise_rng
return pd.DataFrame({"ln_cfu": [1.0]})
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(MagicMock(), {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
return_value=({}, latent_params)), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
side_effect=fake_growth_df):
out_prefix = str(tmp_path / "prior")
@@ -180,20 +180,20 @@ def fake_growth_df(gm, params, draw_idx, noise_rng):
def test_with_noise_passes_rng(self, tmp_path):
"""noise=True must pass a non-None rng to growth_df_from_prior."""
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
captured = {}
latent_params = {"activity": np.array([[1.0]])}
- def fake_growth_df(gm, params, draw_idx, noise_rng):
+ def fake_growth_df(orchestrator, params, draw_idx, noise_rng):
captured["noise_rng"] = noise_rng
return pd.DataFrame({"ln_cfu": [1.0]})
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(MagicMock(), {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
return_value=({}, latent_params)), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
side_effect=fake_growth_df):
out_prefix = str(tmp_path / "prior")
@@ -203,16 +203,16 @@ def fake_growth_df(gm, params, draw_idx, noise_rng):
assert captured["noise_rng"] is not None
def test_csv_content_written(self, tmp_path):
- from tfscreen.tfmodel.scripts.prior_predictive_cli import sample_prior
+ from tfscreen.tfmodel.scripts.sample_prior_cli import sample_prior
latent_params = {"activity": np.array([[1.0]])}
fake_df = pd.DataFrame({"genotype": ["wt", "A1G"], "ln_cfu": [5.0, 4.5]})
- with patch("tfscreen.tfmodel.scripts.prior_predictive_cli.read_configuration",
+ with patch("tfscreen.tfmodel.scripts.sample_prior_cli.read_configuration",
return_value=(MagicMock(), {})), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.draw_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.draw_prior",
return_value=({}, latent_params)), \
- patch("tfscreen.tfmodel.scripts.prior_predictive_cli.growth_df_from_prior",
+ patch("tfscreen.tfmodel.scripts.sample_prior_cli.growth_df_from_prior",
return_value=fake_df):
out_prefix = str(tmp_path / "prior")
diff --git a/tests/tfscreen/tfmodel/scripts/test_subset_genotypes.py b/tests/tfscreen/tfmodel/scripts/test_subset_genotypes_cli.py
similarity index 100%
rename from tests/tfscreen/tfmodel/scripts/test_subset_genotypes.py
rename to tests/tfscreen/tfmodel/scripts/test_subset_genotypes_cli.py
diff --git a/tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py b/tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py
index ad9a4605..f17612c6 100644
--- a/tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_summarize_fit_cli.py
@@ -8,6 +8,7 @@
import json
import os
import warnings
+from pathlib import Path
import matplotlib
matplotlib.use("Agg")
@@ -17,12 +18,25 @@
import pytest
import yaml
+from unittest.mock import MagicMock, patch
+
+# _build_transform_registry registers log_{col} for every sim column, including
+# log_hill_K which has negative values → np.log(negative) = NaN + RuntimeWarning.
+# This is expected behaviour (the log_log_hill_K entry is just NaN), so suppress
+# across this module rather than wrapping every individual test.
+pytestmark = pytest.mark.filterwarnings(
+ "ignore:invalid value encountered in log:RuntimeWarning"
+)
+
from tfscreen.tfmodel.scripts.summarize_fit_cli import (
+ _find_params_or_posterior,
_find_unique,
_json_safe,
_read_all_losses,
_read_final_loss,
_resolve_path,
+ _try_plot_theta_fits,
+ _try_plot_trajectories,
summarize_fit,
)
@@ -69,9 +83,9 @@ def _make_pred_csv(path, n_training=3, extra_genotypes=None):
"genotype": g,
"titrant_name": TITRANT_NAME,
"titrant_conc": tc,
- "median": theta + tc * 0.0001,
- "lower_95": theta - 0.05,
- "upper_95": theta + 0.05,
+ "q0.5": theta + tc * 0.0001,
+ "q0.025": theta - 0.05,
+ "q0.975": theta + 0.05,
"in_training_data": in_train,
})
pd.DataFrame(rows).to_csv(path, index=False)
@@ -112,7 +126,7 @@ def run_dir(tmp_path):
binding_path = str(tmp_path / "binding.csv")
_make_binding_csv(binding_path)
- _make_pred_csv(str(tmp_path / "run_theta_pred.csv"))
+ _make_pred_csv(str(tmp_path / "run_pred_theta.csv"))
_make_guesses_csv(str(tmp_path / "run_guesses.csv"), n_params=25)
_make_losses_txt(str(tmp_path / "run_losses.txt"), final_loss=-4321.0)
_make_config_yaml(
@@ -127,20 +141,25 @@ def run_dir(tmp_path):
return str(tmp_path)
-@pytest.fixture
-def ground_truth_file(tmp_path):
- """Ground-truth file containing the 4th genotype (not in training data)."""
+def _make_ref_theta_csv(path, theta_col="theta_obs"):
+ """Write a ref theta CSV for the 4th genotype using the given column name."""
rows = []
for tc in TITRANT_CONCS:
rows.append({
"genotype": "E4F",
"titrant_name": TITRANT_NAME,
"titrant_conc": tc,
- "theta_obs": THETA_OBS[3] + tc * 0.0001,
+ theta_col: THETA_OBS[3] + tc * 0.0001,
"theta_std": 0.02,
})
- path = str(tmp_path / "ground_truth.csv")
pd.DataFrame(rows).to_csv(path, index=False)
+
+
+@pytest.fixture
+def ref_theta_file(tmp_path):
+ """Ref theta file containing the 4th genotype (not in training data)."""
+ path = str(tmp_path / "ref_theta.csv")
+ _make_ref_theta_csv(path, theta_col="theta_obs")
return path
@@ -321,15 +340,19 @@ class TestSummarizeFitComplete:
def test_json_written(self, run_dir):
summarize_fit(run_dir)
- assert os.path.exists(os.path.join(run_dir, "tfs_summarize_fit_summary.json"))
+ assert os.path.exists(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json"))
+
+ def test_summary_dir_created(self, run_dir):
+ summarize_fit(run_dir)
+ assert os.path.isdir(os.path.join(run_dir, "summary"))
def test_pdf_written(self, run_dir):
summarize_fit(run_dir)
- assert os.path.exists(os.path.join(run_dir, "tfs_summarize_theta_corr.pdf"))
+ assert os.path.exists(os.path.join(run_dir, "summary", "tfs_summarize_theta_corr.pdf"))
def test_json_is_valid(self, run_dir):
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert "metadata" in data
assert "theta" in data
@@ -337,7 +360,7 @@ def test_json_is_valid(self, run_dir):
def test_json_theta_and_growth_have_training_subkey(self, run_dir):
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert "training" in data["theta"]
assert "test" in data["theta"]
@@ -345,7 +368,7 @@ def test_json_theta_and_growth_have_training_subkey(self, run_dir):
def test_metadata_fields_populated(self, run_dir):
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
meta = data["metadata"]
assert meta["n_parameters"] == 25
@@ -354,23 +377,23 @@ def test_metadata_fields_populated(self, run_dir):
def test_theta_training_stats_populated(self, run_dir):
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
stats = data["theta"]["training"]
assert stats is not None
for key in ("rmse", "pearson_r", "spearman_r", "r_squared", "mean_error"):
assert key in stats, f"Missing key: {key}"
- def test_theta_test_null_when_no_ground_truth(self, run_dir):
+ def test_theta_test_null_when_no_ref_theta(self, run_dir):
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert data["theta"]["test"] is None
assert data["metadata"]["n_theta_test_points"] is None
def test_growth_training_null_when_no_growth_pred(self, run_dir):
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert data["growth"]["training"] is None
assert data["metadata"]["n_growth_training_points"] is None
@@ -379,10 +402,10 @@ def test_growth_training_stats_populated_when_file_present(self, run_dir):
pd.DataFrame({
"genotype": ["wt"] * 6,
"ln_cfu": np.linspace(8.0, 13.0, 6),
- "median": np.linspace(8.1, 13.1, 6),
- }).to_csv(os.path.join(run_dir, "tfs_growth_pred.csv"), index=False)
+ "q0.5": np.linspace(8.1, 13.1, 6),
+ }).to_csv(os.path.join(run_dir, "tfs_pred_growth.csv"), index=False)
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
stats = data["growth"]["training"]
assert stats is not None
@@ -394,68 +417,240 @@ def test_growth_training_drops_nan_ln_cfu(self, run_dir):
df = pd.DataFrame({
"genotype": ["wt"] * 6,
"ln_cfu": [8.0, np.nan, 9.0, np.nan, 10.0, 11.0],
- "median": np.linspace(8.1, 13.1, 6),
+ "q0.5": np.linspace(8.1, 13.1, 6),
})
- df.to_csv(os.path.join(run_dir, "tfs_growth_pred.csv"), index=False)
+ df.to_csv(os.path.join(run_dir, "tfs_pred_growth.csv"), index=False)
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert data["metadata"]["n_growth_training_points"] == 4
def test_loss_pdf_written(self, run_dir):
summarize_fit(run_dir)
- assert os.path.exists(os.path.join(run_dir, "tfs_summarize_losses.pdf"))
+ assert os.path.exists(os.path.join(run_dir, "summary", "tfs_summarize_losses.pdf"))
def test_growth_corr_pdf_written_when_file_present(self, run_dir):
# Write a minimal growth_pred CSV
pd.DataFrame({
"genotype": ["wt"] * 6,
"ln_cfu": np.linspace(8.0, 13.0, 6),
- "median": np.linspace(8.1, 13.1, 6),
- }).to_csv(os.path.join(run_dir, "tfs_growth_pred.csv"), index=False)
+ "q0.5": np.linspace(8.1, 13.1, 6),
+ }).to_csv(os.path.join(run_dir, "tfs_pred_growth.csv"), index=False)
summarize_fit(run_dir)
- assert os.path.exists(os.path.join(run_dir, "tfs_summarize_growth_corr.pdf"))
+ assert os.path.exists(os.path.join(run_dir, "summary", "tfs_summarize_growth_corr.pdf"))
def test_growth_corr_pdf_not_written_when_file_absent(self, run_dir):
summarize_fit(run_dir)
- assert not os.path.exists(os.path.join(run_dir, "tfs_summarize_growth_corr.pdf"))
+ assert not os.path.exists(os.path.join(run_dir, "summary", "tfs_summarize_growth_corr.pdf"))
+
+ # ------------------------------------------------------------------
+ # CSV outputs
+ # ------------------------------------------------------------------
+
+ def test_theta_corr_training_csv_written(self, run_dir):
+ summarize_fit(run_dir)
+ csv_path = os.path.join(run_dir, "summary", "tfs_summarize_theta_corr_training.csv")
+ assert os.path.exists(csv_path)
+ df = pd.read_csv(csv_path)
+ assert "ref" in df.columns
+ assert "q0.5" in df.columns
+ assert len(df) > 0
+
+ def test_theta_corr_training_csv_not_written_when_no_binding(self, run_dir):
+ config_path = os.path.join(run_dir, "run_config.yaml")
+ with open(config_path) as fh:
+ cfg = yaml.safe_load(fh)
+ cfg["data"]["binding"] = "/nonexistent/binding.csv"
+ with open(config_path, "w") as fh:
+ yaml.dump(cfg, fh)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ summarize_fit(run_dir)
+ assert not os.path.exists(
+ os.path.join(run_dir, "summary", "tfs_summarize_theta_corr_training.csv")
+ )
+
+ def test_theta_corr_test_csv_not_written_without_ref_theta(self, run_dir):
+ summarize_fit(run_dir)
+ assert not os.path.exists(
+ os.path.join(run_dir, "summary", "tfs_summarize_theta_corr_test.csv")
+ )
+
+ def test_theta_corr_test_csv_written_with_ref_theta_file(self, run_dir, ref_theta_file):
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
+ _make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
+ summarize_fit(run_dir, ref_theta_file=ref_theta_file)
+ csv_path = os.path.join(run_dir, "summary", "tfs_summarize_theta_corr_test.csv")
+ assert os.path.exists(csv_path)
+ df = pd.read_csv(csv_path)
+ assert "ref" in df.columns
+ assert "q0.5" in df.columns
+ assert len(df) == len(TITRANT_CONCS)
+
+ def test_growth_corr_csv_written(self, run_dir):
+ growth_pred_path = os.path.join(run_dir, "tfs_pred_growth.csv")
+ pd.DataFrame({
+ "genotype": ["wt"] * 6,
+ "ln_cfu": np.linspace(8.0, 13.0, 6),
+ "ln_cfu_std": np.ones(6) * 0.1,
+ "q0.5": np.linspace(8.1, 13.1, 6),
+ }).to_csv(growth_pred_path, index=False)
+ summarize_fit(run_dir)
+ csv_path = os.path.join(run_dir, "summary", "tfs_summarize_growth_corr.csv")
+ assert os.path.exists(csv_path)
+ assert not os.path.islink(csv_path)
+ df = pd.read_csv(csv_path)
+ assert "ref" in df.columns
+ assert "ref_std" in df.columns
+ assert "ln_cfu" not in df.columns
+ assert "ln_cfu_std" not in df.columns
+
+ def test_growth_corr_csv_not_written_when_file_absent(self, run_dir):
+ summarize_fit(run_dir)
+ assert not os.path.exists(
+ os.path.join(run_dir, "summary", "tfs_summarize_growth_corr.csv")
+ )
+
+ def test_growth_corr_csv_overwritten_on_rerun(self, run_dir):
+ growth_pred_path = os.path.join(run_dir, "tfs_pred_growth.csv")
+ pd.DataFrame({
+ "genotype": ["wt"] * 3,
+ "ln_cfu": [8.0, 9.0, 10.0],
+ "q0.5": [8.1, 9.1, 10.1],
+ }).to_csv(growth_pred_path, index=False)
+ summarize_fit(run_dir)
+ summarize_fit(run_dir) # second call must not raise
+ csv_path = os.path.join(run_dir, "summary", "tfs_summarize_growth_corr.csv")
+ assert os.path.exists(csv_path)
+ assert not os.path.islink(csv_path)
def test_custom_out_prefix(self, run_dir, tmp_path):
custom_prefix = str(tmp_path / "custom" / "myrun")
- os.makedirs(os.path.dirname(custom_prefix), exist_ok=True)
summarize_fit(run_dir, out_prefix=custom_prefix)
assert os.path.exists(f"{custom_prefix}_fit_summary.json")
assert os.path.exists(f"{custom_prefix}_theta_corr.pdf")
# ---------------------------------------------------------------------------
-# summarize_fit — with ground truth (test statistics)
+# summarize_fit — with ref theta file (test statistics)
# ---------------------------------------------------------------------------
-class TestSummarizeFitWithGroundTruth:
+class TestSummarizeFitWithRefTheta:
- def test_theta_test_stats_populated(self, run_dir, ground_truth_file):
+ def test_theta_test_stats_populated(self, run_dir, ref_theta_file):
# Add out-of-training genotype to pred CSV
- pred_path = os.path.join(run_dir, "run_theta_pred.csv")
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
_make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
- summarize_fit(run_dir, ground_truth_file=ground_truth_file)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ summarize_fit(run_dir, ref_theta_file=ref_theta_file)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert data["theta"]["test"] is not None
assert data["metadata"]["n_theta_test_points"] == len(TITRANT_CONCS)
- def test_theta_test_pearson_r_reasonable(self, run_dir, ground_truth_file):
- pred_path = os.path.join(run_dir, "run_theta_pred.csv")
+ def test_theta_test_pearson_r_reasonable(self, run_dir, ref_theta_file):
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
_make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
- summarize_fit(run_dir, ground_truth_file=ground_truth_file)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ summarize_fit(run_dir, ref_theta_file=ref_theta_file)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
r = data["theta"]["test"]["pearson_r"]
# Predictions are close to truth so r should be high
assert r is not None and r > 0.9
+ def test_theta_col_fallback_to_theta(self, run_dir, tmp_path):
+ """Ref theta file with 'theta' column (no 'theta_obs') still works."""
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
+ _make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
+
+ gt_path = str(tmp_path / "ref_theta_col.csv")
+ _make_ref_theta_csv(gt_path, theta_col="theta")
+
+ summarize_fit(run_dir, ref_theta_file=gt_path)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
+ data = json.load(fh)
+ assert data["theta"]["test"] is not None
+ assert data["metadata"]["n_theta_test_points"] == len(TITRANT_CONCS)
+
+ def test_missing_theta_col_warns_and_skips(self, run_dir, tmp_path):
+ """Ref theta file with no theta column issues a warning and skips test stats."""
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
+ _make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
+
+ gt_path = str(tmp_path / "ref_no_theta.csv")
+ pd.DataFrame({
+ "genotype": ["E4F"],
+ "titrant_name": [TITRANT_NAME],
+ "titrant_conc": [0.0],
+ "some_other_col": [0.5],
+ }).to_csv(gt_path, index=False)
+
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ summarize_fit(run_dir, ref_theta_file=gt_path)
+ assert any("neither" in str(x.message).lower() for x in w)
+
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
+ data = json.load(fh)
+ assert data["theta"]["test"] is None
+ assert data["metadata"]["n_theta_test_points"] is None
+
+ def test_auto_discovers_sim_genotype_theta_in_run_dir(self, run_dir):
+ """*_sim_genotype_theta.csv in run_dir is used automatically when no arg given."""
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
+ _make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
+ _make_ref_theta_csv(
+ os.path.join(run_dir, "tfs_sim_genotype_theta.csv"), theta_col="theta_obs"
+ )
+
+ summarize_fit(run_dir)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
+ data = json.load(fh)
+ assert data["theta"]["test"] is not None
+ assert data["metadata"]["n_theta_test_points"] == len(TITRANT_CONCS)
+ assert data["metadata"]["ref_theta_file"] is not None
+
+ def test_ref_theta_file_arg_overrides_auto_discovery(self, run_dir, tmp_path):
+ """Explicit ref_theta_file takes precedence over auto-discovered file."""
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
+ _make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
+
+ # Auto-discoverable file: has no matching genotypes → would give 0 test points
+ pd.DataFrame({
+ "genotype": ["NOMATCH"],
+ "titrant_name": [TITRANT_NAME],
+ "titrant_conc": [0.0],
+ "theta_obs": [0.5],
+ }).to_csv(os.path.join(run_dir, "tfs_sim_genotype_theta.csv"), index=False)
+
+ # Explicit file: has the real test genotype → should win
+ explicit_path = str(tmp_path / "explicit_ref.csv")
+ _make_ref_theta_csv(explicit_path, theta_col="theta_obs")
+
+ summarize_fit(run_dir, ref_theta_file=explicit_path)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
+ data = json.load(fh)
+ assert data["metadata"]["ref_theta_file"] == explicit_path
+ assert data["metadata"]["n_theta_test_points"] == len(TITRANT_CONCS)
+
+ def test_metadata_ref_theta_file_recorded(self, run_dir, ref_theta_file):
+ """ref_theta_file path is recorded in JSON metadata."""
+ pred_path = os.path.join(run_dir, "run_pred_theta.csv")
+ _make_pred_csv(pred_path, n_training=3, extra_genotypes=["E4F"])
+
+ summarize_fit(run_dir, ref_theta_file=ref_theta_file)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
+ data = json.load(fh)
+ assert data["metadata"]["ref_theta_file"] == ref_theta_file
+
+ def test_metadata_ref_theta_file_null_when_nothing_found(self, run_dir):
+ """ref_theta_file metadata key is None when neither arg nor auto-discovery finds a file."""
+ summarize_fit(run_dir)
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
+ data = json.load(fh)
+ assert data["metadata"]["ref_theta_file"] is None
+
# ---------------------------------------------------------------------------
# summarize_fit — graceful failure cases
@@ -464,11 +659,11 @@ def test_theta_test_pearson_r_reasonable(self, run_dir, ground_truth_file):
class TestSummarizeFitGraceful:
def test_missing_theta_pred_still_writes_json(self, run_dir):
- os.remove(os.path.join(run_dir, "run_theta_pred.csv"))
+ os.remove(os.path.join(run_dir, "run_pred_theta.csv"))
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
summarize_fit(run_dir)
- json_path = os.path.join(run_dir, "tfs_summarize_fit_summary.json")
+ json_path = os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")
assert os.path.exists(json_path)
with open(json_path) as fh:
data = json.load(fh)
@@ -485,7 +680,7 @@ def test_missing_binding_csv_gives_null_theta_training(self, run_dir, tmp_path):
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
summarize_fit(run_dir)
- with open(os.path.join(run_dir, "tfs_summarize_fit_summary.json")) as fh:
+ with open(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json")) as fh:
data = json.load(fh)
assert data["theta"]["training"] is None
@@ -494,7 +689,7 @@ def test_missing_config_still_writes_json(self, run_dir):
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
summarize_fit(run_dir)
- assert os.path.exists(os.path.join(run_dir, "tfs_summarize_fit_summary.json"))
+ assert os.path.exists(os.path.join(run_dir, "summary", "tfs_summarize_fit_summary.json"))
def test_empty_run_dir_writes_json_with_nulls(self, tmp_path):
empty_dir = str(tmp_path / "empty")
@@ -502,7 +697,7 @@ def test_empty_run_dir_writes_json_with_nulls(self, tmp_path):
with warnings.catch_warnings(record=True):
warnings.simplefilter("always")
summarize_fit(empty_dir)
- json_path = os.path.join(empty_dir, "tfs_summarize_fit_summary.json")
+ json_path = os.path.join(empty_dir, "summary", "tfs_summarize_fit_summary.json")
assert os.path.exists(json_path)
with open(json_path) as fh:
data = json.load(fh)
@@ -511,3 +706,1301 @@ def test_empty_run_dir_writes_json_with_nulls(self, tmp_path):
assert data["growth"]["training"] is None
assert data["metadata"]["n_parameters"] is None
assert data["metadata"]["final_loss"] is None
+
+
+# ---------------------------------------------------------------------------
+# _find_params_or_posterior
+# ---------------------------------------------------------------------------
+
+class TestFindParamsOrPosterior:
+
+ def test_returns_none_when_neither_present(self, tmp_path):
+ kind, path = _find_params_or_posterior(str(tmp_path))
+ assert kind is None
+ assert path is None
+
+ def test_finds_posterior_h5(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ kind, path = _find_params_or_posterior(str(tmp_path))
+ assert kind == "posterior"
+ assert path == str(tmp_path / "run_posterior.h5")
+
+ def test_finds_params_npz(self, tmp_path):
+ (tmp_path / "run_params.npz").touch()
+ kind, path = _find_params_or_posterior(str(tmp_path))
+ assert kind == "params"
+ assert path == str(tmp_path / "run_params.npz")
+
+ def test_prefers_posterior_over_params(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ (tmp_path / "run_params.npz").touch()
+ kind, path = _find_params_or_posterior(str(tmp_path))
+ assert kind == "posterior"
+
+ def test_warns_on_multiple_posterior_files(self, tmp_path):
+ (tmp_path / "a_posterior.h5").touch()
+ (tmp_path / "b_posterior.h5").touch()
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ kind, path = _find_params_or_posterior(str(tmp_path))
+ assert kind == "posterior"
+ assert any("posterior" in str(x.message).lower() for x in w)
+
+ def test_warns_on_multiple_params_files_falls_back_to_first(self, tmp_path):
+ (tmp_path / "a_params.npz").touch()
+ (tmp_path / "b_params.npz").touch()
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ kind, path = _find_params_or_posterior(str(tmp_path))
+ assert kind == "params"
+ assert path == str(tmp_path / "a_params.npz")
+ assert any("params" in str(x.message).lower() for x in w)
+
+
+# ---------------------------------------------------------------------------
+# _try_plot_theta_fits
+# ---------------------------------------------------------------------------
+
+_PATCH_PLOT_THETA_FITS = "tfscreen.plot.plot_theta_fits.plot_theta_fits"
+
+_MOCK_BINDING_WITH_STD = pd.DataFrame({
+ "genotype": ["wt", "A1B", "C2D"],
+ "titrant_name": ["IPTG"] * 3,
+ "titrant_conc": [10.0] * 3,
+ "theta_obs": [0.1, 0.5, 0.9],
+ "theta_std": [0.02, 0.02, 0.02],
+})
+
+_MOCK_PRED_FOR_THETA = pd.DataFrame({
+ "genotype": ["wt", "A1B", "C2D"],
+ "titrant_name": ["IPTG"] * 3,
+ "titrant_conc": [10.0] * 3,
+ "q0.5": [0.12, 0.48, 0.88],
+ "in_training_data": [1, 1, 1],
+})
+
+
+class TestTryPlotThetaFits:
+
+ def _make_mock_ax(self):
+ mock_ax = MagicMock()
+ mock_ax.get_figure.return_value = MagicMock()
+ return mock_ax
+
+ # ------------------------------------------------------------------
+ # Guard: None inputs
+ # ------------------------------------------------------------------
+
+ def test_skips_silently_when_binding_df_is_none(self, tmp_path):
+ with patch(_PATCH_PLOT_THETA_FITS) as mock_plot:
+ _try_plot_theta_fits(None, _MOCK_PRED_FOR_THETA, str(tmp_path / "out"))
+ mock_plot.assert_not_called()
+
+ def test_skips_silently_when_pred_df_is_none(self, tmp_path):
+ with patch(_PATCH_PLOT_THETA_FITS) as mock_plot:
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, None, str(tmp_path / "out"))
+ mock_plot.assert_not_called()
+
+ # ------------------------------------------------------------------
+ # Guard: missing theta_std column
+ # ------------------------------------------------------------------
+
+ def test_warns_and_skips_when_theta_std_missing(self, tmp_path):
+ binding_no_std = _MOCK_BINDING_WITH_STD.drop(columns=["theta_std"])
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch(_PATCH_PLOT_THETA_FITS) as mock_plot:
+ _try_plot_theta_fits(binding_no_std, _MOCK_PRED_FOR_THETA,
+ str(tmp_path / "out"))
+ mock_plot.assert_not_called()
+ assert any("theta_std" in str(x.message).lower() for x in w)
+
+ # ------------------------------------------------------------------
+ # Happy path: one call per genotype, savefig called with correct paths
+ # ------------------------------------------------------------------
+
+ def test_calls_plot_once_per_genotype(self, tmp_path):
+ mock_ax = self._make_mock_ax()
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax) as mock_plot:
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, _MOCK_PRED_FOR_THETA,
+ str(tmp_path / "out"))
+ assert mock_plot.call_count == 3
+
+ def test_savefig_called_with_correct_paths(self, tmp_path):
+ out_prefix = str(tmp_path / "out")
+ mock_ax = self._make_mock_ax()
+ mock_fig = mock_ax.get_figure.return_value
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax):
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, _MOCK_PRED_FOR_THETA, out_prefix)
+ saved_paths = {call.args[0] for call in mock_fig.savefig.call_args_list}
+ for geno in ["wt", "A1B", "C2D"]:
+ assert f"{out_prefix}_{geno}_theta_fits.pdf" in saved_paths
+
+ def test_title_set_to_genotype_name(self, tmp_path):
+ mock_ax = self._make_mock_ax()
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax) as mock_plot:
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, _MOCK_PRED_FOR_THETA,
+ str(tmp_path / "out"))
+ titles = {call.kwargs.get("title") or call.args[1]
+ for call in mock_plot.call_args_list
+ if mock_plot.call_args_list}
+ # title is passed as kwarg
+ titles = {call.kwargs["title"] for call in mock_plot.call_args_list}
+ assert titles == {"wt", "A1B", "C2D"}
+
+ # ------------------------------------------------------------------
+ # Filename sanitization
+ # ------------------------------------------------------------------
+
+ def test_slash_in_genotype_sanitized_in_filename(self, tmp_path):
+ out_prefix = str(tmp_path / "out")
+ binding = _MOCK_BINDING_WITH_STD.copy()
+ binding["genotype"] = binding["genotype"].replace("wt", "wt/mut")
+ pred = _MOCK_PRED_FOR_THETA.copy()
+ pred["genotype"] = pred["genotype"].replace("wt", "wt/mut")
+ mock_ax = self._make_mock_ax()
+ mock_fig = mock_ax.get_figure.return_value
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax):
+ _try_plot_theta_fits(binding, pred, out_prefix)
+ saved_paths = {call.args[0] for call in mock_fig.savefig.call_args_list}
+ assert f"{out_prefix}_wt_mut_theta_fits.pdf" in saved_paths
+
+ def test_space_in_genotype_sanitized_in_filename(self, tmp_path):
+ out_prefix = str(tmp_path / "out")
+ binding = _MOCK_BINDING_WITH_STD.copy()
+ binding["genotype"] = binding["genotype"].replace("wt", "wt mut")
+ pred = _MOCK_PRED_FOR_THETA.copy()
+ pred["genotype"] = pred["genotype"].replace("wt", "wt mut")
+ mock_ax = self._make_mock_ax()
+ mock_fig = mock_ax.get_figure.return_value
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax):
+ _try_plot_theta_fits(binding, pred, out_prefix)
+ saved_paths = {call.args[0] for call in mock_fig.savefig.call_args_list}
+ assert f"{out_prefix}_wt_mut_theta_fits.pdf" in saved_paths
+
+ # ------------------------------------------------------------------
+ # Only genotypes present in both dfs are plotted
+ # ------------------------------------------------------------------
+
+ def test_only_genotypes_in_both_dfs_are_plotted(self, tmp_path):
+ pred_subset = _MOCK_PRED_FOR_THETA[_MOCK_PRED_FOR_THETA["genotype"] != "C2D"].copy()
+ mock_ax = self._make_mock_ax()
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax) as mock_plot:
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, pred_subset,
+ str(tmp_path / "out"))
+ assert mock_plot.call_count == 2
+
+ # ------------------------------------------------------------------
+ # CSV output: one CSV per genotype alongside each PDF
+ # ------------------------------------------------------------------
+
+ def test_writes_csv_per_genotype(self, tmp_path):
+ out_prefix = str(tmp_path / "out")
+ mock_ax = self._make_mock_ax()
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax):
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, _MOCK_PRED_FOR_THETA, out_prefix)
+ for geno in ["wt", "A1B", "C2D"]:
+ csv_path = f"{out_prefix}_{geno}_theta_fits.csv"
+ assert os.path.exists(csv_path), f"Missing CSV for {geno}"
+ df = pd.read_csv(csv_path)
+ assert "theta_obs" in df.columns
+ assert "q0.5" in df.columns
+ assert all(df["genotype"] == geno)
+
+ def test_csv_written_before_pdf(self, tmp_path):
+ """CSV must exist even if savefig raises."""
+ out_prefix = str(tmp_path / "out")
+ mock_ax = self._make_mock_ax()
+ mock_ax.get_figure.return_value.savefig.side_effect = RuntimeError("render failed")
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax):
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, _MOCK_PRED_FOR_THETA, out_prefix)
+ # savefig raised inside the try block so warn should fire, but the
+ # CSV for the first genotype (A1B, alphabetically first) must exist
+ # because it is written before savefig is called
+ assert os.path.exists(f"{out_prefix}_A1B_theta_fits.csv")
+
+ # ------------------------------------------------------------------
+ # Guard: plot raises
+ # ------------------------------------------------------------------
+
+ def test_warns_gracefully_when_plot_raises(self, tmp_path):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch(_PATCH_PLOT_THETA_FITS, side_effect=RuntimeError("render failed")):
+ _try_plot_theta_fits(_MOCK_BINDING_WITH_STD, _MOCK_PRED_FOR_THETA,
+ str(tmp_path / "out"))
+ assert any("theta fit" in str(x.message).lower() for x in w)
+
+ # ------------------------------------------------------------------
+ # Integration: summarize_fit calls _try_plot_theta_fits
+ # ------------------------------------------------------------------
+
+ def test_summarize_fit_calls_plot_theta_fits_per_training_genotype(self, run_dir):
+ mock_ax = MagicMock()
+ mock_ax.get_figure.return_value = MagicMock()
+ with patch(_PATCH_PLOT_THETA_FITS, return_value=mock_ax) as mock_plot:
+ summarize_fit(run_dir)
+ # run_dir fixture has 3 training genotypes
+ assert mock_plot.call_count == 3
+
+ def test_summarize_fit_skips_theta_fits_when_no_pred_csv(self, run_dir):
+ os.remove(os.path.join(run_dir, "run_pred_theta.csv"))
+ with patch(_PATCH_PLOT_THETA_FITS) as mock_plot:
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ summarize_fit(run_dir)
+ mock_plot.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# _try_plot_trajectories
+# ---------------------------------------------------------------------------
+
+_MOCK_BINDING_DF = pd.DataFrame({
+ "genotype": ["wt", "A1B", "C2D"],
+ "titrant_name": ["IPTG"] * 3,
+ "titrant_conc": [0.0] * 3,
+ "theta_obs": [0.1, 0.5, 0.9],
+})
+
+_CONFIG_WITH_GROWTH = {"data": {"growth": "growth.csv", "binding": "binding.csv"}}
+_CONFIG_NO_GROWTH = {"data": {"binding": "binding.csv"}}
+
+_PATCH_PREDICT_DF = "tfscreen.plot.geno_trajectory.predict_geno_trajectory_df"
+_PATCH_PLOT_GENO = "tfscreen.plot.geno_trajectory.plot_geno_trajectory"
+_PATCH_READ_CONFIG = "tfscreen.tfmodel.configuration_io.read_configuration"
+
+# DataFrame returned by the mock predict_geno_trajectory_df — one row per
+# genotype so the per-genotype loop fires exactly len(GENOTYPES) times.
+_MOCK_PRED_DF = pd.DataFrame({
+ "genotype": list(_MOCK_BINDING_DF["genotype"]),
+})
+
+
+class TestTryPlotTrajectories:
+
+ # ------------------------------------------------------------------
+ # Guard: no growth data in config
+ # ------------------------------------------------------------------
+
+ def test_skips_silently_when_no_growth_key(self, tmp_path):
+ with patch(_PATCH_PREDICT_DF) as mock_pdf:
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_NO_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ mock_pdf.assert_not_called()
+
+ def test_skips_silently_when_config_yaml_is_none(self, tmp_path):
+ with patch(_PATCH_PREDICT_DF) as mock_pdf:
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=None,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ mock_pdf.assert_not_called()
+
+ # ------------------------------------------------------------------
+ # Guard: no params / posterior file
+ # ------------------------------------------------------------------
+
+ def test_warns_and_skips_when_no_pred_file(self, tmp_path):
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch(_PATCH_PREDICT_DF) as mock_pdf:
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ mock_pdf.assert_not_called()
+ assert any("posterior" in str(x.message).lower() or
+ "params" in str(x.message).lower() for x in w)
+
+ # ------------------------------------------------------------------
+ # Guard: read_configuration raises
+ # ------------------------------------------------------------------
+
+ def test_warns_and_skips_when_orchestrator_load_fails(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch(_PATCH_READ_CONFIG, side_effect=RuntimeError("load failed")), \
+ patch(_PATCH_PREDICT_DF) as mock_pdf:
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ mock_pdf.assert_not_called()
+ assert any("trajectory" in str(x.message).lower() for x in w)
+
+ # ------------------------------------------------------------------
+ # Happy path: posterior.h5 present — passes h5 path as second arg
+ # ------------------------------------------------------------------
+
+ def test_calls_predict_df_with_h5_path(self, tmp_path):
+ h5_path = str(tmp_path / "run_posterior.h5")
+ (tmp_path / "run_posterior.h5").touch()
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF) as mock_pdf, \
+ patch(_PATCH_PLOT_GENO, return_value=MagicMock()):
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ mock_pdf.assert_called_once()
+ args, _ = mock_pdf.call_args
+ assert args[1] == h5_path
+
+ # ------------------------------------------------------------------
+ # Happy path: only params.npz present (fallback) — passes npz path
+ # ------------------------------------------------------------------
+
+ def test_calls_predict_df_with_npz_path(self, tmp_path):
+ npz_path = str(tmp_path / "run_params.npz")
+ (tmp_path / "run_params.npz").touch()
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF) as mock_pdf, \
+ patch(_PATCH_PLOT_GENO, return_value=MagicMock()):
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ mock_pdf.assert_called_once()
+ args, _ = mock_pdf.call_args
+ assert args[1] == npz_path
+
+ # ------------------------------------------------------------------
+ # Happy path: one PDF saved per genotype
+ # ------------------------------------------------------------------
+
+ def test_saves_per_genotype_pdfs(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ out_prefix = str(tmp_path / "out")
+ mock_fig = MagicMock()
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF), \
+ patch(_PATCH_PLOT_GENO, return_value=mock_fig) as mock_plot:
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=out_prefix,
+ binding_df=_MOCK_BINDING_DF,
+ )
+ genotypes = sorted(_MOCK_PRED_DF["genotype"].unique().tolist(), key=str)
+ assert mock_plot.call_count == len(genotypes)
+ saved_paths = {call.args[0] for call in mock_fig.savefig.call_args_list}
+ for geno in genotypes:
+ safe = geno.replace("/", "_").replace(" ", "_")
+ assert f"{out_prefix}_{safe}_trajectory.pdf" in saved_paths
+
+ # ------------------------------------------------------------------
+ # Genotype list derived from binding_df
+ # ------------------------------------------------------------------
+
+ def test_passes_genotypes_from_binding_df(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF) as mock_pdf, \
+ patch(_PATCH_PLOT_GENO, return_value=MagicMock()):
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ _, kwargs = mock_pdf.call_args
+ assert set(kwargs["genotypes"]) == {"wt", "A1B", "C2D"}
+
+ def test_passes_none_genotypes_when_binding_df_is_none(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF) as mock_pdf, \
+ patch(_PATCH_PLOT_GENO, return_value=MagicMock()):
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=None,
+ )
+ _, kwargs = mock_pdf.call_args
+ assert kwargs["genotypes"] is None
+
+ # ------------------------------------------------------------------
+ # CSV output: one CSV per genotype alongside each trajectory PDF
+ # ------------------------------------------------------------------
+
+ def test_writes_csv_per_genotype(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ out_prefix = str(tmp_path / "out")
+ mock_fig = MagicMock()
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF), \
+ patch(_PATCH_PLOT_GENO, return_value=mock_fig):
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=out_prefix,
+ binding_df=_MOCK_BINDING_DF,
+ )
+ for geno in sorted(_MOCK_PRED_DF["genotype"].unique().tolist(), key=str):
+ safe = geno.replace("/", "_").replace(" ", "_")
+ csv_path = f"{out_prefix}_{safe}_trajectory.csv"
+ assert os.path.exists(csv_path), f"Missing CSV for {geno}"
+ df = pd.read_csv(csv_path)
+ assert all(df["genotype"] == geno)
+
+ # ------------------------------------------------------------------
+ # Guard: predict_geno_trajectory_df raises
+ # ------------------------------------------------------------------
+
+ def test_warns_gracefully_when_plot_raises(self, tmp_path):
+ (tmp_path / "run_posterior.h5").touch()
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, side_effect=RuntimeError("plot failed")):
+ _try_plot_trajectories(
+ config_file=str(tmp_path / "cfg.yaml"),
+ config_yaml=_CONFIG_WITH_GROWTH,
+ run_dir=str(tmp_path),
+ out_prefix=str(tmp_path / "out"),
+ binding_df=_MOCK_BINDING_DF,
+ )
+ assert any("trajectory" in str(x.message).lower() for x in w)
+
+ # ------------------------------------------------------------------
+ # Integration: summarize_fit calls _try_plot_trajectories
+ # ------------------------------------------------------------------
+
+ def test_summarize_fit_calls_predict_df_when_h5_and_growth_config_present(
+ self, run_dir
+ ):
+ config_path = os.path.join(run_dir, "run_config.yaml")
+ with open(config_path) as fh:
+ cfg = yaml.safe_load(fh)
+ cfg["data"]["growth"] = "growth.csv"
+ with open(config_path, "w") as fh:
+ yaml.dump(cfg, fh)
+ (Path(run_dir) / "run_posterior.h5").touch()
+
+ with patch(_PATCH_READ_CONFIG, return_value=(MagicMock(), {})), \
+ patch(_PATCH_PREDICT_DF, return_value=_MOCK_PRED_DF) as mock_pdf, \
+ patch(_PATCH_PLOT_GENO, return_value=MagicMock()):
+ summarize_fit(run_dir)
+
+ mock_pdf.assert_called_once()
+
+ def test_summarize_fit_skips_plot_when_no_pred_file(self, run_dir):
+ config_path = os.path.join(run_dir, "run_config.yaml")
+ with open(config_path) as fh:
+ cfg = yaml.safe_load(fh)
+ cfg["data"]["growth"] = "growth.csv"
+ with open(config_path, "w") as fh:
+ yaml.dump(cfg, fh)
+ # No posterior.h5 or params.npz written.
+
+ with patch(_PATCH_PREDICT_DF) as mock_pdf:
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ summarize_fit(run_dir)
+
+ mock_pdf.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# _build_transform_registry / _compute_*_obs / _summarize_params
+# ---------------------------------------------------------------------------
+
+from tfscreen.tfmodel.scripts.summarize_fit_cli import (
+ _build_transform_registry,
+ _compute_direct_obs,
+ _compute_diff_obs,
+ _compute_epi_obs,
+ _extract_quantile_data,
+ _summarize_params,
+ _try_calibration,
+ _try_plot_params_corr,
+)
+
+# Shared sim_parameters DataFrame used across param-summary tests.
+_SIM_DF = pd.DataFrame([
+ {"genotype": "wt", "activity": 1.0, "theta_low": 0.90, "theta_high": 0.05, "log_hill_K": -4.0},
+ {"genotype": "A1G", "activity": 0.5, "theta_low": 0.70, "theta_high": 0.08, "log_hill_K": -3.5},
+ {"genotype": "B2H", "activity": 2.0, "theta_low": 0.80, "theta_high": 0.03, "log_hill_K": -4.5},
+ {"genotype": "C3K", "activity": 1.5, "theta_low": 0.85, "theta_high": 0.04, "log_hill_K": -4.1},
+ {"genotype": "A1G/B2H", "activity": 1.2, "theta_low": 0.75, "theta_high": 0.06, "log_hill_K": -4.2},
+ {"genotype": "A1G/C3K", "activity": 0.8, "theta_low": 0.78, "theta_high": 0.07, "log_hill_K": -3.9},
+])
+
+
+def _logit(x):
+ x = np.clip(x, 1e-9, 1 - 1e-9)
+ return np.log(x / (1 - x))
+
+
+class TestBuildTransformRegistry:
+
+ def test_direct_column_registered(self):
+ reg = _build_transform_registry(_SIM_DF)
+ assert "activity" in reg
+ assert reg["activity"]["wt"] == pytest.approx(1.0)
+ assert reg["activity"]["A1G"] == pytest.approx(0.5)
+
+ def test_log_prefix_registered(self):
+ reg = _build_transform_registry(_SIM_DF)
+ assert "log_activity" in reg
+ assert reg["log_activity"]["A1G"] == pytest.approx(np.log(0.5))
+
+ def test_logit_theta_col_registered_under_long_name(self):
+ reg = _build_transform_registry(_SIM_DF)
+ assert "logit_theta_low" in reg
+ assert reg["logit_theta_low"]["wt"] == pytest.approx(_logit(0.90))
+
+ def test_logit_theta_col_also_registered_under_short_name(self):
+ # theta_low → logit_low (short form used by hill_mut component)
+ reg = _build_transform_registry(_SIM_DF)
+ assert "logit_low" in reg
+ assert reg["logit_low"]["A1G"] == pytest.approx(_logit(0.70))
+
+ def test_logit_delta_compound_registered(self):
+ reg = _build_transform_registry(_SIM_DF)
+ expected_wt = _logit(0.05) - _logit(0.90)
+ assert "logit_delta" in reg
+ assert reg["logit_delta"]["wt"] == pytest.approx(expected_wt)
+
+ def test_logit_delta_absent_when_theta_high_missing(self):
+ df = _SIM_DF.drop(columns=["theta_high"])
+ reg = _build_transform_registry(df)
+ assert "logit_delta" not in reg
+
+ def test_direct_column_log_hill_K_not_double_logged(self):
+ # log_hill_K is already logged in sim; should appear as direct match.
+ reg = _build_transform_registry(_SIM_DF)
+ assert "log_hill_K" in reg
+ assert reg["log_hill_K"]["wt"] == pytest.approx(-4.0)
+
+
+class TestComputeDirectObs:
+
+ def test_returns_obs_for_genotype_key(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"genotype": ["wt", "A1G", "B2H"], "q0.5": [1.0, 0.4, 1.9]})
+ obs = _compute_direct_obs(params_df, obs_s)
+ assert obs == pytest.approx([1.0, 0.5, 2.0])
+
+ def test_returns_none_when_no_genotype_column(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"condition": ["A"], "q0.5": [1.0]})
+ assert _compute_direct_obs(params_df, obs_s) is None
+
+ def test_nan_for_unknown_genotype(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"genotype": ["wt", "UNKNOWN"], "q0.5": [1.0, 0.5]})
+ obs = _compute_direct_obs(params_df, obs_s)
+ assert obs[0] == pytest.approx(1.0)
+ assert np.isnan(obs[1])
+
+ def test_extra_columns_preserved(self):
+ # The function only computes obs; it does not modify params_df
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({
+ "genotype": ["wt"],
+ "titrant_name": ["iptg"],
+ "q0.5": [0.9],
+ })
+ obs = _compute_direct_obs(params_df, obs_s)
+ assert obs == pytest.approx([1.0])
+
+
+class TestComputeDiffObs:
+
+ def test_computes_diff_from_wt(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ # activity: wt=1.0, A1G=0.5, B2H=2.0
+ params_df = pd.DataFrame({
+ "titrant_name": ["iptg", "iptg"],
+ "mutation": ["A1G", "B2H"],
+ "q0.5": [-0.55, 0.95],
+ })
+ obs = _compute_diff_obs(params_df, obs_s)
+ assert obs == pytest.approx([-0.5, 1.0])
+
+ def test_returns_none_when_no_mutation_column(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"genotype": ["A1G"], "q0.5": [0.5]})
+ assert _compute_diff_obs(params_df, obs_s) is None
+
+ def test_returns_nan_when_wt_absent(self):
+ df_no_wt = _SIM_DF[_SIM_DF["genotype"] != "wt"].copy()
+ obs_s = df_no_wt.set_index("genotype")["activity"]
+ params_df = pd.DataFrame({"mutation": ["A1G"], "q0.5": [0.4]})
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ obs = _compute_diff_obs(params_df, obs_s)
+ assert np.all(np.isnan(obs))
+ assert any("wt" in str(x.message).lower() for x in w)
+
+ def test_nan_for_unknown_mutation(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"mutation": ["A1G", "UNKNOWN"], "q0.5": [0.0, 0.0]})
+ obs = _compute_diff_obs(params_df, obs_s)
+ assert obs[0] == pytest.approx(-0.5)
+ assert np.isnan(obs[1])
+
+ def test_diff_with_logit_low_transform(self):
+ obs_s = _build_transform_registry(_SIM_DF)["logit_low"]
+ # logit(theta_low[A1G]) - logit(theta_low[wt]) = logit(0.7) - logit(0.9)
+ params_df = pd.DataFrame({"mutation": ["A1G"], "q0.5": [0.0]})
+ obs = _compute_diff_obs(params_df, obs_s)
+ expected = _logit(0.70) - _logit(0.90)
+ assert obs == pytest.approx([expected])
+
+
+class TestComputeEpiObs:
+
+ def test_computes_epistasis_for_double_mutant(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ # ep = (activity[A1G/B2H] - activity[B2H]) - (activity[A1G] - activity[wt])
+ # = (1.2 - 2.0) - (0.5 - 1.0) = -0.8 - (-0.5) = -0.3
+ params_df = pd.DataFrame({"pair": ["A1G/B2H"], "q0.5": [0.0]})
+ obs = _compute_epi_obs(params_df, obs_s, "activity")
+ assert obs == pytest.approx([-0.3])
+
+ def test_returns_none_when_no_pair_column(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"genotype": ["A1G/B2H"], "q0.5": [0.0]})
+ assert _compute_epi_obs(params_df, obs_s, "activity") is None
+
+ def test_nan_for_unknown_pair(self):
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"pair": ["A1G/B2H", "X1Y/Z2W"], "q0.5": [0.0, 0.0]})
+ obs = _compute_epi_obs(params_df, obs_s, "activity")
+ assert obs[0] == pytest.approx(-0.3)
+ assert np.isnan(obs[1])
+
+ def test_reversed_pair_order_resolved_via_standardize(self):
+ # extract_epistasis standardizes 'B2H/A1G' → 'A1G/B2H'; the pair
+ # column in the params file may use the opposite order.
+ obs_s = _build_transform_registry(_SIM_DF)["activity"]
+ params_df = pd.DataFrame({"pair": ["B2H/A1G"], "q0.5": [0.0]})
+ obs = _compute_epi_obs(params_df, obs_s, "activity")
+ assert obs == pytest.approx([-0.3])
+
+ def test_logit_delta_compound_epi(self):
+ obs_s = _build_transform_registry(_SIM_DF)["logit_delta"]
+ params_df = pd.DataFrame({"pair": ["A1G/B2H"], "q0.5": [0.0]})
+ obs = _compute_epi_obs(params_df, obs_s, "logit_delta")
+ # Manually compute expected epistasis
+ def ld(row): return _logit(row["theta_high"]) - _logit(row["theta_low"])
+ sim_idx = _SIM_DF.set_index("genotype")
+ ld_wt = ld(sim_idx.loc["wt"])
+ ld_a1g = ld(sim_idx.loc["A1G"])
+ ld_b2h = ld(sim_idx.loc["B2H"])
+ ld_dbl = ld(sim_idx.loc["A1G/B2H"])
+ expected = (ld_dbl - ld_b2h) - (ld_a1g - ld_wt)
+ assert obs == pytest.approx([expected])
+
+
+class TestTryPlotParamsCorr:
+
+ def _make_df(self, n=10, nan_obs=False):
+ obs = np.linspace(0.0, 1.0, n)
+ med = obs + np.random.default_rng(0).normal(0, 0.05, n)
+ df = pd.DataFrame({"ref": obs, "q0.5": med})
+ if nan_obs:
+ df.loc[::3, "ref"] = np.nan
+ return df
+
+ def test_pdf_written(self, tmp_path):
+ pdf = str(tmp_path / "out.pdf")
+ _try_plot_params_corr(self._make_df(), "activity", pdf)
+ assert os.path.exists(pdf)
+
+ def test_skips_when_fewer_than_two_valid_rows(self, tmp_path):
+ pdf = str(tmp_path / "out.pdf")
+ df = pd.DataFrame({"ref": [np.nan, np.nan], "q0.5": [0.5, 0.6]})
+ _try_plot_params_corr(df, "activity", pdf)
+ assert not os.path.exists(pdf)
+
+ def test_skips_when_only_one_valid_row(self, tmp_path):
+ pdf = str(tmp_path / "out.pdf")
+ df = pd.DataFrame({"ref": [1.0, np.nan], "q0.5": [1.1, 0.5]})
+ _try_plot_params_corr(df, "activity", pdf)
+ assert not os.path.exists(pdf)
+
+ def test_nan_rows_dropped_before_plot(self, tmp_path):
+ pdf = str(tmp_path / "out.pdf")
+ _try_plot_params_corr(self._make_df(nan_obs=True), "activity", pdf)
+ assert os.path.exists(pdf)
+
+ def test_warns_gracefully_on_failure(self, tmp_path):
+ pdf = str(tmp_path / "out.pdf")
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch("tfscreen.tfmodel.scripts.summarize_fit_cli.xy_corr",
+ side_effect=RuntimeError("render failed")):
+ _try_plot_params_corr(self._make_df(), "activity", pdf)
+ assert not os.path.exists(pdf)
+ assert any("correlation plot" in str(x.message).lower() for x in w)
+
+ def test_label_used_in_axis_titles(self, tmp_path):
+ """Axis labels include the param name; just check no crash with unusual label."""
+ pdf = str(tmp_path / "out.pdf")
+ _try_plot_params_corr(self._make_df(), "logit_low", pdf)
+ assert os.path.exists(pdf)
+
+
+class TestSummarizeParams:
+ """Integration tests for _summarize_params using real file I/O."""
+
+ def _write_sim_params(self, path):
+ _SIM_DF.to_csv(path, index=False)
+
+ def _write_direct_params(self, path, col="activity"):
+ # genotype + titrant_name + median columns (as tfs_params_activity.csv looks)
+ rows = []
+ for g in _SIM_DF["genotype"]:
+ rows.append({"genotype": g, "titrant_name": "iptg", "q0.5": 0.0})
+ pd.DataFrame(rows).to_csv(path, index=False)
+
+ def _write_diff_params(self, path, col="activity"):
+ rows = []
+ for g in ["A1G", "B2H"]:
+ rows.append({"titrant_name": "iptg", "mutation": g, "q0.5": 0.0})
+ pd.DataFrame(rows).to_csv(path, index=False)
+
+ def _write_epi_params(self, path, col="activity"):
+ rows = [
+ {"titrant_name": "iptg", "pair": "A1G/B2H", "q0.5": 0.0},
+ {"titrant_name": "iptg", "pair": "A1G/C3K", "q0.5": 0.0},
+ ]
+ pd.DataFrame(rows).to_csv(path, index=False)
+
+ def test_no_output_when_sim_params_absent(self, tmp_path):
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ _summarize_params(str(tmp_path), out_prefix)
+ assert not any(f.name.startswith("tfs_summarize_params")
+ for f in (tmp_path / "summary").iterdir())
+
+ def test_direct_csv_written_with_obs(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_direct_params(str(tmp_path / "run_params_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ out_path = out_prefix + "_params_activity.csv"
+ assert os.path.exists(out_path)
+ df = pd.read_csv(out_path)
+ assert "ref" in df.columns
+ expected = dict(zip(_SIM_DF["genotype"], _SIM_DF["activity"]))
+ for _, row in df.iterrows():
+ assert row["ref"] == pytest.approx(expected[row["genotype"]])
+
+ def test_diff_csv_written_with_obs(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_diff_params(str(tmp_path / "run_params_d_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ out_path = out_prefix + "_params_d_activity.csv"
+ assert os.path.exists(out_path)
+ df = pd.read_csv(out_path)
+ assert "ref" in df.columns
+ wt_act = 1.0
+ for _, row in df.iterrows():
+ act = {"A1G": 0.5, "B2H": 2.0}[row["mutation"]]
+ assert row["ref"] == pytest.approx(act - wt_act)
+
+ def test_epi_csv_written_with_obs(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_epi_params(str(tmp_path / "run_params_epi_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ out_path = out_prefix + "_params_epi_activity.csv"
+ assert os.path.exists(out_path)
+ df = pd.read_csv(out_path)
+ assert "ref" in df.columns
+ # ep = (1.2 - 2.0) - (0.5 - 1.0) = -0.3
+ assert df.loc[df["pair"] == "A1G/B2H", "ref"].values[0] == pytest.approx(-0.3)
+
+ def test_logit_low_diff_obs(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ pd.DataFrame([
+ {"titrant_name": "iptg", "mutation": "A1G", "q0.5": 0.0},
+ ]).to_csv(str(tmp_path / "run_params_d_logit_low.csv"), index=False)
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ out_path = out_prefix + "_params_d_logit_low.csv"
+ assert os.path.exists(out_path)
+ df = pd.read_csv(out_path)
+ expected = _logit(0.70) - _logit(0.90)
+ assert df["ref"].values[0] == pytest.approx(expected)
+
+ def test_logit_delta_epi_obs(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ pd.DataFrame([
+ {"titrant_name": "iptg", "pair": "A1G/B2H", "q0.5": 0.0},
+ ]).to_csv(str(tmp_path / "run_params_epi_logit_delta.csv"), index=False)
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ out_path = out_prefix + "_params_epi_logit_delta.csv"
+ assert os.path.exists(out_path)
+ df = pd.read_csv(out_path)
+ sim_idx = _SIM_DF.set_index("genotype")
+ def ld(g): return _logit(sim_idx.loc[g, "theta_high"]) - _logit(sim_idx.loc[g, "theta_low"])
+ expected = (ld("A1G/B2H") - ld("B2H")) - (ld("A1G") - ld("wt"))
+ assert df["ref"].values[0] == pytest.approx(expected)
+
+ def test_warns_when_expected_direct_file_missing(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ # Deliberately omit *_params_activity.csv
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ assert any("activity" in str(x.message) for x in w)
+
+ def test_output_written_to_summary_dir(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_direct_params(str(tmp_path / "run_params_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ summary_dir = tmp_path / "summary"
+ written = [f.name for f in summary_dir.iterdir()]
+ assert any("params_activity" in n for n in written)
+
+ def test_direct_pdf_written_alongside_csv(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_direct_params(str(tmp_path / "run_params_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ assert os.path.exists(out_prefix + "_params_activity.pdf")
+
+ def test_diff_pdf_written_alongside_csv(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_diff_params(str(tmp_path / "run_params_d_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ assert os.path.exists(out_prefix + "_params_d_activity.pdf")
+
+ def test_epi_pdf_written_alongside_csv(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ self._write_epi_params(str(tmp_path / "run_params_epi_activity.csv"))
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ assert os.path.exists(out_prefix + "_params_epi_activity.pdf")
+
+ def test_file_without_genotype_key_skipped_silently(self, tmp_path):
+ self._write_sim_params(str(tmp_path / "run_sim_parameters.csv"))
+ # ln_cfu0-style file: no genotype column
+ pd.DataFrame([
+ {"condition_rep": "kanR+kan", "q0.5": 30.5},
+ ]).to_csv(str(tmp_path / "run_params_ln_cfu0.csv"), index=False)
+ out_prefix = str(tmp_path / "summary" / "tfs_summarize")
+ os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ _summarize_params(str(tmp_path), out_prefix)
+ # No crash and no output file for ln_cfu0 (not in sim_params columns)
+ summary_dir = tmp_path / "summary"
+ assert not any("ln_cfu0" in f.name for f in summary_dir.iterdir())
+
+ def test_summarize_fit_calls_summarize_params(self, run_dir):
+ """_summarize_params is invoked from summarize_fit end-to-end."""
+ sim_path = os.path.join(run_dir, "tfs_sim_parameters.csv")
+ _SIM_DF.to_csv(sim_path, index=False)
+ params_path = os.path.join(run_dir, "tfs_params_activity.csv")
+ rows = [{"genotype": g, "titrant_name": "iptg", "q0.5": 0.0}
+ for g in _SIM_DF["genotype"]]
+ pd.DataFrame(rows).to_csv(params_path, index=False)
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ summarize_fit(run_dir)
+ out_csv = os.path.join(run_dir, "summary", "tfs_summarize_params_activity.csv")
+ out_pdf = os.path.join(run_dir, "summary", "tfs_summarize_params_activity.pdf")
+ assert os.path.exists(out_csv)
+ assert os.path.exists(out_pdf)
+ df = pd.read_csv(out_csv)
+ assert "ref" in df.columns
+
+
+# ---------------------------------------------------------------------------
+# _extract_quantile_data
+# ---------------------------------------------------------------------------
+
+_PATCH_CALIBRATION_SUMMARY = (
+ "tfscreen.tfmodel.scripts.summarize_fit_cli.calibration_summary"
+)
+
+
+class TestExtractQuantileData:
+
+ def test_standard_q_cols_returns_matrix_and_levels(self):
+ df = pd.DataFrame({
+ "q0.025": [0.1, 0.2],
+ "q0.5": [0.5, 0.6],
+ "q0.975": [0.9, 1.0],
+ "ref": [0.4, 0.5],
+ })
+ mat, levels = _extract_quantile_data(df)
+ assert levels == pytest.approx([0.025, 0.5, 0.975])
+ assert mat.shape == (2, 3)
+ assert mat[0, 1] == pytest.approx(0.5) # q0.5 for first row
+
+ def test_levels_sorted_numerically_not_lexicographically(self):
+ # q0.1 < q0.25 numerically; lexicographic would put q0.25 < q0.1
+ df = pd.DataFrame({"q0.25": [0.3], "q0.1": [0.1], "q0.9": [0.8]})
+ _, levels = _extract_quantile_data(df)
+ assert list(levels) == pytest.approx([0.1, 0.25, 0.9])
+
+ def test_returns_none_none_when_no_q_cols(self):
+ df = pd.DataFrame({"ref": [0.1], "genotype": ["wt"]})
+ mat, levels = _extract_quantile_data(df)
+ assert mat is None
+ assert levels is None
+
+ def test_returns_none_none_with_exactly_one_q_col(self):
+ df = pd.DataFrame({"q0.5": [0.5], "ref": [0.4]})
+ mat, levels = _extract_quantile_data(df)
+ assert mat is None
+ assert levels is None
+
+ def test_non_q_columns_excluded(self):
+ df = pd.DataFrame({
+ "q0.5": [0.5],
+ "q0.975": [0.9],
+ "quality": [1.0], # must not match
+ })
+ mat, levels = _extract_quantile_data(df)
+ assert mat is not None
+ assert mat.shape == (1, 2)
+ assert list(levels) == pytest.approx([0.5, 0.975])
+
+ def test_returns_values_in_same_row_order_as_df(self):
+ df = pd.DataFrame({
+ "q0.025": [0.10, 0.20, 0.30],
+ "q0.975": [0.80, 0.85, 0.90],
+ })
+ mat, _ = _extract_quantile_data(df)
+ assert mat[1, 0] == pytest.approx(0.20)
+ assert mat[2, 1] == pytest.approx(0.90)
+
+
+# ---------------------------------------------------------------------------
+# _try_calibration
+# ---------------------------------------------------------------------------
+
+def _make_calibration_df(n=20, q_cols=True, ref_col="ref"):
+ """Build a minimal DataFrame for calibration tests."""
+ rng = np.random.default_rng(42)
+ data = {}
+ if ref_col:
+ data[ref_col] = rng.uniform(0, 1, n)
+ if q_cols:
+ data["q0.025"] = rng.uniform(0.0, 0.3, n)
+ data["q0.5"] = rng.uniform(0.3, 0.7, n)
+ data["q0.975"] = rng.uniform(0.7, 1.0, n)
+ else:
+ data["q0.5"] = rng.uniform(0.3, 0.7, n)
+ return pd.DataFrame(data)
+
+
+class TestTryCalibration:
+
+ def test_calls_calibration_summary_with_correct_shapes(self, tmp_path):
+ df = _make_calibration_df(n=30)
+ out_prefix = str(tmp_path / "cal")
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ _try_calibration(df, "ref", out_prefix, "test label")
+ mock_cal.assert_called_once()
+ _, kwargs = mock_cal.call_args
+ assert kwargs["true_vals"].shape == (30,)
+ assert kwargs["quantile_matrix"].shape == (30, 3)
+ assert kwargs["quantile_levels"].tolist() == pytest.approx([0.025, 0.5, 0.975])
+ assert kwargs["out_prefix"] == out_prefix
+ assert kwargs["label"] == "test label"
+
+ def test_skips_silently_when_df_is_none(self, tmp_path):
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ _try_calibration(None, "ref", str(tmp_path / "cal"), "label")
+ mock_cal.assert_not_called()
+
+ def test_skips_silently_when_ref_col_absent(self, tmp_path):
+ df = _make_calibration_df()
+ df = df.drop(columns=["ref"])
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ _try_calibration(df, "ref", str(tmp_path / "cal"), "label")
+ mock_cal.assert_not_called()
+
+ def test_skips_silently_when_fewer_than_two_q_cols(self, tmp_path):
+ df = _make_calibration_df(q_cols=False) # only q0.5
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ _try_calibration(df, "ref", str(tmp_path / "cal"), "label")
+ mock_cal.assert_not_called()
+
+ def test_warns_gracefully_when_calibration_summary_raises(self, tmp_path):
+ df = _make_calibration_df()
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ with patch(_PATCH_CALIBRATION_SUMMARY, side_effect=RuntimeError("boom")):
+ _try_calibration(df, "ref", str(tmp_path / "cal"), "my label")
+ assert any("my label" in str(x.message) for x in w)
+
+ def test_passes_ref_col_values_as_true_vals(self, tmp_path):
+ rng = np.random.default_rng(7)
+ ref_vals = rng.uniform(0, 1, 10)
+ df = pd.DataFrame({
+ "ref": ref_vals,
+ "q0.1": rng.uniform(0, 0.4, 10),
+ "q0.9": rng.uniform(0.6, 1.0, 10),
+ })
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ _try_calibration(df, "ref", str(tmp_path / "cal"), "label")
+ _, kwargs = mock_cal.call_args
+ np.testing.assert_array_equal(kwargs["true_vals"], ref_vals)
+
+ def test_works_with_non_default_ref_col_name(self, tmp_path):
+ df = _make_calibration_df(ref_col="ln_cfu")
+ df = df.rename(columns={"ref": "other"}) # ensure "ref" is absent
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ _try_calibration(df, "ln_cfu", str(tmp_path / "cal"), "growth")
+ mock_cal.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# Integration: calibration called from each summarize_fit call site
+# ---------------------------------------------------------------------------
+
+def _make_pred_csv_with_quantiles(path, n_training=3, extra_genotypes=None):
+ """Like _make_pred_csv but includes q0.025 / q0.5 / q0.975."""
+ rows = []
+ genotypes = list(GENOTYPES[:n_training])
+ if extra_genotypes:
+ genotypes += extra_genotypes
+ theta_map = dict(zip(GENOTYPES, THETA_PRED))
+ for g in genotypes:
+ in_train = 1 if g in GENOTYPES[:n_training] else 0
+ theta = theta_map.get(g, 0.5)
+ for tc in TITRANT_CONCS:
+ rows.append({
+ "genotype": g,
+ "titrant_name": TITRANT_NAME,
+ "titrant_conc": tc,
+ "q0.025": max(0.0, theta - 0.10),
+ "q0.5": theta + tc * 0.0001,
+ "q0.975": min(1.0, theta + 0.10),
+ "in_training_data": in_train,
+ })
+ pd.DataFrame(rows).to_csv(path, index=False)
+
+
+class TestCalibrationIntegration:
+ """Verify calibration_summary is (or isn't) called from each wiring point."""
+
+ # ------------------------------------------------------------------
+ # Theta training
+ # ------------------------------------------------------------------
+
+ def test_calibration_called_for_theta_training_with_quantile_cols(self, run_dir):
+ _make_pred_csv_with_quantiles(os.path.join(run_dir, "run_pred_theta.csv"))
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert any("theta_training" in p for p in prefixes)
+
+ def test_calibration_skipped_for_theta_training_without_quantile_cols(self, run_dir):
+ # Overwrite with a pred CSV that has only q0.5 (no interval columns)
+ rows = []
+ for g in GENOTYPES[:3]:
+ for tc in TITRANT_CONCS:
+ rows.append({
+ "genotype": g, "titrant_name": TITRANT_NAME,
+ "titrant_conc": tc, "q0.5": 0.5, "in_training_data": 1,
+ })
+ pd.DataFrame(rows).to_csv(
+ os.path.join(run_dir, "run_pred_theta.csv"), index=False
+ )
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert not any("theta_training" in p for p in prefixes)
+
+ # ------------------------------------------------------------------
+ # Theta test
+ # ------------------------------------------------------------------
+
+ def test_calibration_called_for_theta_test_with_quantile_cols(
+ self, run_dir, ref_theta_file
+ ):
+ _make_pred_csv_with_quantiles(
+ os.path.join(run_dir, "run_pred_theta.csv"),
+ n_training=3, extra_genotypes=["E4F"],
+ )
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir, ref_theta_file=ref_theta_file)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert any("theta_test" in p for p in prefixes)
+
+ def test_calibration_skipped_for_theta_test_when_no_ref_theta_file(self, run_dir):
+ _make_pred_csv_with_quantiles(os.path.join(run_dir, "run_pred_theta.csv"))
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert not any("theta_test" in p for p in prefixes)
+
+ # ------------------------------------------------------------------
+ # Growth
+ # ------------------------------------------------------------------
+
+ def test_calibration_called_for_growth_with_quantile_cols(self, run_dir):
+ pd.DataFrame({
+ "genotype": ["wt"] * 6,
+ "ln_cfu": np.linspace(8.0, 13.0, 6),
+ "q0.025": np.linspace(7.5, 12.5, 6),
+ "q0.5": np.linspace(8.0, 13.0, 6),
+ "q0.975": np.linspace(8.5, 13.5, 6),
+ }).to_csv(os.path.join(run_dir, "tfs_pred_growth.csv"), index=False)
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert any("growth" in p for p in prefixes)
+
+ def test_calibration_uses_ref_col_for_growth(self, run_dir):
+ pd.DataFrame({
+ "genotype": ["wt"] * 6,
+ "ln_cfu": np.linspace(8.0, 13.0, 6),
+ "q0.025": np.linspace(7.5, 12.5, 6),
+ "q0.5": np.linspace(8.0, 13.0, 6),
+ "q0.975": np.linspace(8.5, 13.5, 6),
+ }).to_csv(os.path.join(run_dir, "tfs_pred_growth.csv"), index=False)
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ growth_calls = [c for c in mock_cal.call_args_list
+ if "growth" in c.kwargs["out_prefix"]]
+ assert len(growth_calls) == 1
+ # true_vals must come from the renamed "ref" column (was ln_cfu)
+ np.testing.assert_array_almost_equal(
+ growth_calls[0].kwargs["true_vals"],
+ np.linspace(8.0, 13.0, 6),
+ )
+
+ def test_calibration_skipped_for_growth_without_quantile_cols(self, run_dir):
+ pd.DataFrame({
+ "genotype": ["wt"] * 6,
+ "ln_cfu": np.linspace(8.0, 13.0, 6),
+ "q0.5": np.linspace(8.0, 13.0, 6),
+ }).to_csv(os.path.join(run_dir, "tfs_pred_growth.csv"), index=False)
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert not any("growth" in p for p in prefixes)
+
+ def test_calibration_skipped_for_growth_when_no_growth_pred_file(self, run_dir):
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert not any("growth" in p for p in prefixes)
+
+ # ------------------------------------------------------------------
+ # Params
+ # ------------------------------------------------------------------
+
+ def test_calibration_called_for_params_with_quantile_cols(self, run_dir):
+ _SIM_DF.to_csv(os.path.join(run_dir, "tfs_sim_parameters.csv"), index=False)
+ rows = [
+ {"genotype": g, "titrant_name": "iptg",
+ "q0.025": 0.0, "q0.5": 0.0, "q0.975": 0.0}
+ for g in _SIM_DF["genotype"]
+ ]
+ pd.DataFrame(rows).to_csv(
+ os.path.join(run_dir, "tfs_params_activity.csv"), index=False
+ )
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert any("params_activity" in p for p in prefixes)
+
+ def test_calibration_skipped_for_params_without_quantile_cols(self, run_dir):
+ _SIM_DF.to_csv(os.path.join(run_dir, "tfs_sim_parameters.csv"), index=False)
+ rows = [{"genotype": g, "titrant_name": "iptg", "q0.5": 0.0}
+ for g in _SIM_DF["genotype"]]
+ pd.DataFrame(rows).to_csv(
+ os.path.join(run_dir, "tfs_params_activity.csv"), index=False
+ )
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ prefixes = [call.kwargs["out_prefix"] for call in mock_cal.call_args_list]
+ assert not any("params_activity" in p for p in prefixes)
+
+ def test_params_calibration_out_prefix_strips_csv_extension(self, run_dir):
+ _SIM_DF.to_csv(os.path.join(run_dir, "tfs_sim_parameters.csv"), index=False)
+ rows = [
+ {"genotype": g, "titrant_name": "iptg",
+ "q0.025": 0.0, "q0.5": 0.0, "q0.975": 0.0}
+ for g in _SIM_DF["genotype"]
+ ]
+ pd.DataFrame(rows).to_csv(
+ os.path.join(run_dir, "tfs_params_activity.csv"), index=False
+ )
+ with warnings.catch_warnings(record=True):
+ warnings.simplefilter("always")
+ with patch(_PATCH_CALIBRATION_SUMMARY) as mock_cal:
+ summarize_fit(run_dir)
+ params_calls = [c for c in mock_cal.call_args_list
+ if "params_activity" in c.kwargs["out_prefix"]]
+ assert len(params_calls) == 1
+ assert not params_calls[0].kwargs["out_prefix"].endswith(".csv")
diff --git a/tests/tfscreen/tfmodel/scripts/test_summarize_sbc_cli.py b/tests/tfscreen/tfmodel/scripts/test_summarize_sbc_cli.py
index 874a2a71..7d974d4c 100644
--- a/tests/tfscreen/tfmodel/scripts/test_summarize_sbc_cli.py
+++ b/tests/tfscreen/tfmodel/scripts/test_summarize_sbc_cli.py
@@ -15,5 +15,5 @@ def test_main_calls_generalized_main_with_summarize_sbc(self):
def test_summarize_sbc_imported_from_analysis(self):
from tfscreen.tfmodel.scripts.summarize_sbc_cli import summarize_sbc
- from tfscreen.tfmodel.analysis.sbc import summarize_sbc as canonical
+ from tfscreen.tfmodel.analysis.error_calibration import summarize_sbc as canonical
assert summarize_sbc is canonical
diff --git a/tests/tfscreen/tfmodel/test_checkpoint_protection.py b/tests/tfscreen/tfmodel/test_checkpoint_protection.py
index 4b621b3b..a3d9cc84 100644
--- a/tests/tfscreen/tfmodel/test_checkpoint_protection.py
+++ b/tests/tfscreen/tfmodel/test_checkpoint_protection.py
@@ -5,11 +5,11 @@
@pytest.fixture
def mock_growth_model(mocker):
- mock_gm_class = mocker.patch("tfscreen.tfmodel.scripts.fit_model_cli.read_configuration")
- mock_gm_instance = MagicMock()
+ mock_orchestrator_class = mocker.patch("tfscreen.tfmodel.scripts.fit_model_cli.read_configuration")
+ mock_orchestrator_instance = MagicMock()
mock_init_params = {"a": 1.0}
- mock_gm_class.return_value = (mock_gm_instance, mock_init_params)
- return mock_gm_class, mock_gm_instance
+ mock_orchestrator_class.return_value = (mock_orchestrator_instance, mock_init_params)
+ return mock_orchestrator_class, mock_orchestrator_instance
@pytest.fixture
def mock_run_inference(mocker):
diff --git a/tests/tfscreen/tfmodel/test_configure_and_run.py b/tests/tfscreen/tfmodel/test_configure_and_run.py
index b48a0604..1d1a5269 100644
--- a/tests/tfscreen/tfmodel/test_configure_and_run.py
+++ b/tests/tfscreen/tfmodel/test_configure_and_run.py
@@ -14,24 +14,25 @@
from tfscreen.tfmodel.scripts.fit_model_cli import fit_model
from tfscreen.tfmodel.configuration_io import (
read_configuration,
- _update_dataclass
+ _update_dataclass,
+ TFMODEL_KNOWN_KEYS,
)
@pytest.fixture
-def mock_gm(mocker):
- mock_gm_class = mocker.patch("tfscreen.tfmodel.scripts.configure_model_cli.ModelOrchestrator")
- mock_gm_inst = mock_gm_class.return_value
+def mock_orchestrator(mocker):
+ mock_orchestrator_class = mocker.patch("tfscreen.tfmodel.scripts.configure_model_cli.ModelOrchestrator")
+ mock_orchestrator_inst = mock_orchestrator_class.return_value
# Mock settings
- mock_gm_inst.settings = {"batch_size": 128, "theta": "hill_geno"}
+ mock_orchestrator_inst.settings = {"batch_size": 128, "theta": "hill_geno"}
# Mock priors
mock_priors = MagicMock()
mock_priors.val = 1.0
- mock_gm_inst.priors = mock_priors
+ mock_orchestrator_inst.priors = mock_priors
# Mock init_params
- mock_gm_inst.init_params = {
+ mock_orchestrator_inst.init_params = {
"scalar_p": 0.5,
"arr_p": np.array([1.0, 2.0]),
"cond_p": np.array([0.1]),
@@ -48,12 +49,12 @@ def mock_gm(mocker):
"theta": pd.DataFrame({"titrant_name": ["T1"], "titrant_conc": [0.1], "genotype": ["G1"], "map_theta": [0]}),
"ln_cfu0": pd.DataFrame({"replicate": ["R1"], "condition_pre": ["CP1"], "genotype": ["G1"], "map_ln_cfu0": [0]})
}
- mock_gm_inst.growth_tm = mock_tm
+ mock_orchestrator_inst.growth_tm = mock_tm
- return mock_gm_class, mock_gm_inst
+ return mock_orchestrator_class, mock_orchestrator_inst
-def test_configure_growth_analysis_coverage(mock_gm, tmpdir):
- _, mock_gm_inst = mock_gm
+def test_configure_growth_analysis_coverage(mock_orchestrator, tmpdir):
+ _, mock_orchestrator_inst = mock_orchestrator
out_prefix = os.path.join(tmpdir, "test")
# Mock priors with dict and dataclass members for coverage
@@ -65,10 +66,10 @@ def test_configure_growth_analysis_coverage(mock_gm, tmpdir):
mock_priors.dc = mock_sub
mock_priors.d = {"k": 2.0, "arr": np.zeros((2,))}
mock_priors.scalar = 3.0
- mock_gm_inst.priors = mock_priors
+ mock_orchestrator_inst.priors = mock_priors
# Add some specific parameter names to trigger branch coverage
- mock_gm_inst.init_params = {
+ mock_orchestrator_inst.init_params = {
"scalar": 1.0,
"condition_growth_offset": np.array([0.5]),
"theta_logit_low": np.array([0.6]),
@@ -115,17 +116,17 @@ def test_read_configuration_logic(tmpdir, mocker):
"flat_index": [0, 0, 0]
}).to_csv(guesses_path, index=False)
- mock_gm_class = mocker.patch("tfscreen.tfmodel.configuration_io.ModelOrchestrator")
- mock_gm_inst = mock_gm_class.return_value
- mock_gm_inst.init_params = {"param1": 0.0, "param2": jnp.zeros((1,)), "param3": 0.0}
+ mock_orchestrator_class = mocker.patch("tfscreen.tfmodel.configuration_io.ModelOrchestrator")
+ mock_orchestrator_inst = mock_orchestrator_class.return_value
+ mock_orchestrator_inst.init_params = {"param1": 0.0, "param2": jnp.zeros((1,)), "param3": 0.0}
# Mock update_dataclass
mocker.patch("tfscreen.tfmodel.configuration_io._update_dataclass"
, return_value=MagicMock())
- gm, init_params = read_configuration(config_path)
-
- assert gm == mock_gm_inst
+ orchestrator, init_params = read_configuration(config_path)
+
+ assert orchestrator == mock_orchestrator_inst
assert init_params["param1"] == 0.5
assert isinstance(init_params["param2"], jnp.ndarray)
@@ -163,9 +164,9 @@ def test_read_configuration_errors(tmpdir):
with open(config_path, "w") as f:
yaml.dump({"data":{"growth":"g", "binding":"b"}, "components":{}, "priors_file": "priors.csv", "guesses_file": "guesses.csv"}, f)
- with patch("tfscreen.tfmodel.configuration_io.ModelOrchestrator") as mock_gm_class:
- mock_gm_inst = mock_gm_class.return_value
- mock_gm_inst.init_params = {"param1": 0, "param2": 0} # param2 missing in CSV
+ with patch("tfscreen.tfmodel.configuration_io.ModelOrchestrator") as mock_orchestrator_class:
+ mock_orchestrator_inst = mock_orchestrator_class.return_value
+ mock_orchestrator_inst.init_params = {"param1": 0, "param2": 0} # param2 missing in CSV
with pytest.raises(ValueError, match="Missing initial guesses for parameters"):
read_configuration(config_path)
@@ -175,6 +176,64 @@ def test_read_configuration_errors(tmpdir):
with pytest.raises(FileNotFoundError, match="Guesses file not found"):
read_configuration(config_path)
+ # 7. Stale guesses file: array parameter has wrong number of values
+ guesses_path2 = os.path.join(tmpdir, "guesses_stale.csv")
+ pd.DataFrame({
+ "parameter": ["param_arr", "param_arr"],
+ "value": [1.0, 2.0],
+ "flat_index": [0, 1]
+ }).to_csv(guesses_path2, index=False)
+ with open(config_path, "w") as f:
+ yaml.dump({"data":{"growth":"g", "binding":"b"}, "components":{},
+ "priors_file": "priors.csv", "guesses_file": "guesses_stale.csv"}, f)
+ with patch("tfscreen.tfmodel.configuration_io.ModelOrchestrator") as mock_orch_cls:
+ mock_orch_inst = mock_orch_cls.return_value
+ # Model expects shape (5,) but CSV only has 2 values
+ mock_orch_inst.init_params = {"param_arr": jnp.zeros((5,))}
+ with patch("tfscreen.tfmodel.configuration_io._update_dataclass",
+ return_value=MagicMock()):
+ with pytest.raises(ValueError, match="stale"):
+ read_configuration(config_path)
+
+
+# ----------------------------------------------------------------------------
+# test TFMODEL_KNOWN_KEYS / unknown-key validation in read_configuration
+# ----------------------------------------------------------------------------
+
+def test_tfmodel_known_keys_is_frozenset():
+ assert isinstance(TFMODEL_KNOWN_KEYS, frozenset)
+ assert "data" in TFMODEL_KNOWN_KEYS
+ assert "components" in TFMODEL_KNOWN_KEYS
+
+
+def test_read_configuration_unknown_key_raises(tmpdir):
+ config_path = os.path.join(tmpdir, "bad_config.yaml")
+ with open(config_path, "w") as f:
+ yaml.dump({
+ "data": {"growth": "g"},
+ "components": {},
+ "priors_file": "priors.csv",
+ "guesses_file": "guesses.csv",
+ "not_a_real_key": 42,
+ }, f)
+ with pytest.raises(ValueError, match="not_a_real_key"):
+ read_configuration(config_path)
+
+
+def test_read_configuration_unknown_key_error_mentions_label(tmpdir):
+ config_path = os.path.join(tmpdir, "bad_config.yaml")
+ with open(config_path, "w") as f:
+ yaml.dump({
+ "data": {"growth": "g"},
+ "components": {},
+ "priors_file": "priors.csv",
+ "guesses_file": "guesses.csv",
+ "typo_key": "oops",
+ }, f)
+ with pytest.raises(ValueError, match="tfmodel config"):
+ read_configuration(config_path)
+
+
def test_configure_model_errors(tmpdir):
# Missing binding_df
with pytest.raises(ValueError, match="binding_df must be provided"):
@@ -332,11 +391,11 @@ def test_fit_model_svi_full(tmpdir, mocker):
with patch("os.path.exists", return_value=False):
config_path = os.path.join(tmpdir, "config.yaml")
# Minimal valid setup for read_configuration
- mock_gm_run = MagicMock()
- mock_gm_run.settings = {"batch_size": 256}
- mock_gm_run._ln_cfu_df = "g.csv"
- mock_gm_run._binding_df = "b.csv"
- mocker.patch("tfscreen.tfmodel.scripts.fit_model_cli.read_configuration", return_value=(mock_gm_run, {"p":1.0}))
+ mock_orchestrator_run = MagicMock()
+ mock_orchestrator_run.settings = {"batch_size": 256}
+ mock_orchestrator_run._ln_cfu_df = "g.csv"
+ mock_orchestrator_run._binding_df = "b.csv"
+ mocker.patch("tfscreen.tfmodel.scripts.fit_model_cli.read_configuration", return_value=(mock_orchestrator_run, {"p":1.0}))
# Mock write_configuration to avoid yaml problems with mocks
mock_ri_class = mocker.patch("tfscreen.tfmodel.scripts.fit_model_cli.RunInference")
@@ -405,7 +464,7 @@ def test_incompatible_cg_tr_values_are_strings(self):
assert isinstance(tr, str)
assert isinstance(reason, str) and len(reason) > 0
- def test_configure_model_raises_on_incompatible(self, mock_gm):
+ def test_configure_model_raises_on_incompatible(self, mock_orchestrator):
"""configure_model must call check_component_compatibility
before building the model."""
with pytest.raises(ValueError, match="Incompatible model components"):
diff --git a/tests/tfscreen/tfmodel/test_extract_growth_predictions.py b/tests/tfscreen/tfmodel/test_extract_growth_predictions.py
index 80ed5162..9fa473d4 100644
--- a/tests/tfscreen/tfmodel/test_extract_growth_predictions.py
+++ b/tests/tfscreen/tfmodel/test_extract_growth_predictions.py
@@ -56,30 +56,27 @@ def mock_posteriors():
def test_extract_growth_predictions_basic(mock_model, mock_posteriors):
"""Test default behavior for extracting growth predictions."""
results = extract_growth_predictions(mock_model, mock_posteriors)
-
+
# Check output structure
assert isinstance(results, pd.DataFrame)
- assert "median" in results.columns
+ assert "q0.5" in results.columns
assert len(results) == len(mock_model.growth_df)
-
+
# Verify values
- # Row 0 -> median should be 10.5
- assert results.iloc[0]["median"] == 10.5
- # Row 1 -> median should be 11.5
- assert results.iloc[1]["median"] == 11.5
+ # Row 0 -> q0.5 should be 10.5
+ assert results.iloc[0]["q0.5"] == 10.5
+ # Row 1 -> q0.5 should be 11.5
+ assert results.iloc[1]["q0.5"] == 11.5
def test_extract_growth_predictions_custom_quantiles(mock_model, mock_posteriors):
"""Test extracting custom quantiles."""
- q_to_get = {"mean": 0.5, "low": 0.1, "high": 0.9}
- # Add columns to mock_model.growth_df for this test
- for q in q_to_get:
- mock_model.growth_df[q] = np.nan
- results = extract_growth_predictions(mock_model, mock_posteriors, q_to_get=q_to_get)
-
- assert "mean" in results.columns
- assert "low" in results.columns
- assert "high" in results.columns
- assert results.iloc[0]["mean"] == 10.5
+ results = extract_growth_predictions(mock_model, mock_posteriors,
+ q_to_get=[0.1, 0.5, 0.9])
+
+ assert "q0.5" in results.columns
+ assert "q0.1" in results.columns
+ assert "q0.9" in results.columns
+ assert results.iloc[0]["q0.5"] == 10.5
def test_extract_growth_predictions_missing_field(mock_model):
"""Test error handling when growth_pred is missing from posteriors."""
@@ -88,8 +85,8 @@ def test_extract_growth_predictions_missing_field(mock_model):
def test_extract_growth_predictions_invalid_quantiles(mock_model, mock_posteriors):
"""Test error handling for bad quantile input."""
- with pytest.raises(ValueError, match="q_to_get should be a dictionary"):
- extract_growth_predictions(mock_model, mock_posteriors, q_to_get=[0.5])
+ with pytest.raises(ValueError, match="should be a 1-D array-like"):
+ extract_growth_predictions(mock_model, mock_posteriors, q_to_get={"x": 0.5})
def test_extract_growth_predictions_file_loading(mock_model):
"""Test loading posteriors from file."""
diff --git a/tests/tfscreen/tfmodel/test_extract_parameters_congression.py b/tests/tfscreen/tfmodel/test_extract_parameters_congression.py
index 04db7fd8..4e6522a5 100644
--- a/tests/tfscreen/tfmodel/test_extract_parameters_congression.py
+++ b/tests/tfscreen/tfmodel/test_extract_parameters_congression.py
@@ -59,18 +59,18 @@ def test_extract_parameters_congression(mock_model_congression, mock_posteriors_
lam_df = params["lam"]
assert len(lam_df) == 1
assert lam_df.iloc[0]["parameter"] == "lam"
- assert lam_df.iloc[0]["median"] == 1.2
+ assert lam_df.iloc[0]["q0.5"] == 1.2
# Check mu
mu_df = params["mu"]
assert len(mu_df) == 2
assert set(mu_df["titrant_conc"]) == {0.0, 1.0}
- assert np.allclose(mu_df["median"], 0.5)
+ assert np.allclose(mu_df["q0.5"], 0.5)
# Check sigma
sigma_df = params["sigma"]
assert len(sigma_df) == 2
- assert np.allclose(sigma_df["median"], 0.1)
+ assert np.allclose(sigma_df["q0.5"], 0.1)
def test_extract_parameters_no_congression(mock_model_congression):
"""Test that congression parameters are NOT extracted when transformation is none."""
diff --git a/tests/tfscreen/tfmodel/test_extract_samples.py b/tests/tfscreen/tfmodel/test_extract_samples.py
index 49e9ab1b..661cf802 100644
--- a/tests/tfscreen/tfmodel/test_extract_samples.py
+++ b/tests/tfscreen/tfmodel/test_extract_samples.py
@@ -23,7 +23,7 @@
# Shared helpers
# ---------------------------------------------------------------------------
-_Q = {"median": 0.5} # single quantile keeps assertions simple
+_Q = [0.5] # single quantile keeps assertions simple
def _hill_model():
@@ -158,7 +158,7 @@ def test_num_samples_adds_correct_number_of_columns(self):
def test_quantile_columns_still_present_alongside_samples(self):
model = _hill_model()
df = extract_theta_curves(model, _hill_posteriors(), q_to_get=_Q, num_samples=3)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
def test_sample_values_within_posterior_range(self):
"""Sample columns must be actual draws, not summaries outside [min, max]."""
@@ -185,7 +185,7 @@ def test_constant_posteriors_samples_equal_median(self):
df = extract_theta_curves(model, _hill_posteriors(S=10), q_to_get=_Q, num_samples=4)
for i in range(4):
np.testing.assert_allclose(df[f"sample_{i}"].values,
- df["median"].values,
+ df["q0.5"].values,
rtol=1e-6)
def test_num_samples_exceeds_posterior_size_uses_replacement(self):
@@ -233,7 +233,7 @@ def test_sample_columns_are_distinct_rows_of_theta_samples(self):
sample_cols = [c for c in df.columns if c.startswith("sample_")]
assert len(sample_cols) == 5
# At least one sample column should differ from the median somewhere
- diffs = [not np.allclose(df[c].values, df["median"].values) for c in sample_cols]
+ diffs = [not np.allclose(df[c].values, df["q0.5"].values) for c in sample_cols]
assert any(diffs), "All sample columns equal the median — draws are not varying"
@@ -254,7 +254,7 @@ def test_constant_posteriors_samples_equal_median(self):
df = extract_theta_curves(model, _hill_mut_posteriors(S=8), q_to_get=_Q, num_samples=3)
for i in range(3):
np.testing.assert_allclose(df[f"sample_{i}"].values,
- df["median"].values,
+ df["q0.5"].values,
rtol=1e-6)
def test_index_columns_dropped(self):
@@ -289,7 +289,7 @@ def test_constant_posteriors_samples_equal_median(self):
num_samples=3)
for i in range(3):
np.testing.assert_allclose(df[f"sample_{i}"].values,
- df["median"].values,
+ df["q0.5"].values,
rtol=1e-5)
def test_index_columns_dropped(self):
@@ -339,7 +339,7 @@ def test_quantile_columns_still_present(self):
model = _tf_model()
df = extract_growth_predictions(model, _growth_posteriors(), q_to_get=_Q,
num_samples=3)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
def test_constant_posteriors_samples_equal_median(self):
"""When all posterior draws are the same value, samples == median."""
@@ -349,7 +349,7 @@ def test_constant_posteriors_samples_equal_median(self):
df = extract_growth_predictions(model, posteriors, q_to_get=_Q, num_samples=4)
for i in range(4):
np.testing.assert_allclose(df[f"sample_{i}"].values,
- df["median"].values,
+ df["q0.5"].values,
rtol=1e-6)
def test_sample_values_within_posterior_range(self):
diff --git a/tests/tfscreen/tfmodel/test_extract_theta_curves.py b/tests/tfscreen/tfmodel/test_extract_theta_curves.py
index 6d46448f..3183e657 100644
--- a/tests/tfscreen/tfmodel/test_extract_theta_curves.py
+++ b/tests/tfscreen/tfmodel/test_extract_theta_curves.py
@@ -49,18 +49,18 @@ def test_extract_theta_curves_basic(mock_model, mock_posteriors):
assert "genotype" in results.columns
assert "titrant_name" in results.columns
assert "titrant_conc" in results.columns
- assert "median" in results.columns
-
+ assert "q0.5" in results.columns
+
# Should have unique (genotype, titrant_name, titrant_conc)
# wt: 0.0, 1.0; mut: 0.0, 1.0 -> 4 rows
assert len(results) == 4
-
+
# Verify values for wt at conc=1.0
# hill_n=2, log_K=-1.0 (K=0.367), high=0.9, low=0.1
# occupancy = 1 / (1 + exp(-2 * (log(1.0) - (-1.0)))) = 1 / (1 + exp(-2)) = 1 / (1 + 0.135) = 0.88
# theta = 0.1 + (0.9 - 0.1) * 0.88 = 0.1 + 0.8 * 0.88 = 0.804
wt_1 = results[(results["genotype"] == "wt") & (results["titrant_conc"] == 1.0)]
- assert np.allclose(wt_1["median"], 0.804, atol=1e-3)
+ assert np.allclose(wt_1["q0.5"], 0.804, atol=1e-3)
def test_extract_theta_curves_manual_df(mock_model, mock_posteriors):
"""Test providing a manual titrant DataFrame."""
@@ -155,14 +155,14 @@ def _make_unmeasured_model(theta_name="hill_mut"):
def _fake_predict_unmeasured(target_genotypes, titrant_names,
manual_titrant_df, mut_labels,
pair_labels, param_posteriors, q_to_get):
- """Returns one row per (genotype, titrant_conc) with median=0.5."""
+ """Returns one row per (genotype, titrant_conc) with q0.5=0.5."""
rows = []
for g in target_genotypes:
for _, row in manual_titrant_df.iterrows():
rows.append({"genotype": g,
"titrant_name": row["titrant_name"],
"titrant_conc": row["titrant_conc"],
- "median": 0.5})
+ "q0.5": 0.5})
return pd.DataFrame(rows)
@@ -177,7 +177,7 @@ def patched_unmeasured_module():
{"theta": registry_patch},
), patch(
"tfscreen.tfmodel.analysis.extraction.load_posteriors",
- return_value=({"median": 0.5}, {}),
+ return_value=({"q0.5": 0.5}, {}),
):
yield fake_module
@@ -257,7 +257,7 @@ def test_missing_predict_unmeasured_raises(self):
{"theta": {}},
), patch(
"tfscreen.tfmodel.analysis.extraction.load_posteriors",
- return_value=({"median": 0.5}, {}),
+ return_value=({"q0.5": 0.5}, {}),
):
with pytest.raises(ValueError, match="predict_unmeasured"):
extract_theta_unmeasured(
diff --git a/tests/tfscreen/tfmodel/test_extraction_new_models.py b/tests/tfscreen/tfmodel/test_extraction_new_models.py
index 5cae8d8d..b25c70bb 100644
--- a/tests/tfscreen/tfmodel/test_extraction_new_models.py
+++ b/tests/tfscreen/tfmodel/test_extraction_new_models.py
@@ -97,7 +97,7 @@ def _make_model(theta="none", dk_geno="none", activity="fixed",
return model
-_Q = {"median": 0.5} # single quantile to keep assertions simple
+_Q = [0.5] # single quantile to keep assertions simple
# ---------------------------------------------------------------------------
@@ -135,7 +135,7 @@ def test_assembled_output_shape_genotype_x_titrant(self):
# 3 genotypes × 2 titrant names = 6 rows
df = params["theta_low"]
assert len(df) == 6
- assert set(df.columns) >= {"genotype", "titrant_name", "median"}
+ assert set(df.columns) >= {"genotype", "titrant_name", "q0.5"}
def test_mutation_output_shape_mutation_x_titrant(self):
model = _make_model(theta="hill_mut")
@@ -143,7 +143,7 @@ def test_mutation_output_shape_mutation_x_titrant(self):
# 2 mutations × 2 titrant names = 4 rows
df = params["d_logit_low"]
assert len(df) == 4
- assert set(df.columns) >= {"mutation", "titrant_name", "median"}
+ assert set(df.columns) >= {"mutation", "titrant_name", "q0.5"}
def test_compound_index_selects_correct_posterior(self):
"""Each row should use the posterior slice for its (titrant_name_idx, genotype_idx)."""
@@ -165,21 +165,21 @@ def test_compound_index_selects_correct_posterior(self):
# --- assembled (T, G) indexing ---
df = params["theta_low"].set_index(["titrant_name", "genotype"])
- assert df.loc[("iptg", "wt"), "median"] == pytest.approx(0.0)
- assert df.loc[("iptg", "B"), "median"] == pytest.approx(2.0)
- assert df.loc[("atc", "A"), "median"] == pytest.approx(4.0)
- assert df.loc[("atc", "B"), "median"] == pytest.approx(5.0)
+ assert df.loc[("iptg", "wt"), "q0.5"] == pytest.approx(0.0)
+ assert df.loc[("iptg", "B"), "q0.5"] == pytest.approx(2.0)
+ assert df.loc[("atc", "A"), "q0.5"] == pytest.approx(4.0)
+ assert df.loc[("atc", "B"), "q0.5"] == pytest.approx(5.0)
# --- per-mutation (T, M) indexing: flat index = titrant_idx*M + mut_idx ---
dd = params["d_logit_low"].set_index(["titrant_name", "mutation"])
# (iptg=0, A=0) → flat 0*2+0=0 → value 0.0
- assert dd.loc[("iptg", "A"), "median"] == pytest.approx(0.0)
+ assert dd.loc[("iptg", "A"), "q0.5"] == pytest.approx(0.0)
# (iptg=0, B=1) → flat 0*2+1=1 → value 1.0
- assert dd.loc[("iptg", "B"), "median"] == pytest.approx(1.0)
+ assert dd.loc[("iptg", "B"), "q0.5"] == pytest.approx(1.0)
# (atc=1, A=0) → flat 1*2+0=2 → value 2.0
- assert dd.loc[("atc", "A"), "median"] == pytest.approx(2.0)
+ assert dd.loc[("atc", "A"), "q0.5"] == pytest.approx(2.0)
# (atc=1, B=1) → flat 1*2+1=3 → value 3.0
- assert dd.loc[("atc", "B"), "median"] == pytest.approx(3.0)
+ assert dd.loc[("atc", "B"), "q0.5"] == pytest.approx(3.0)
# ---------------------------------------------------------------------------
@@ -342,7 +342,7 @@ def test_default_no_manual_df_output_shape(self, model):
assert "genotype" in df.columns
assert "titrant_name" in df.columns
assert "titrant_conc" in df.columns
- assert "median" in df.columns
+ assert "q0.5" in df.columns
def test_index_columns_dropped_from_output(self, model):
df = extract_theta_curves(model, self._flat_posteriors(),
@@ -373,7 +373,7 @@ def test_hill_equation_math(self, model):
q_to_get=_Q, manual_titrant_df=manual_df, num_samples=None)
expected_occ = 1.0 / (1.0 + np.exp(-2.0 * (0.0 - (-1.0))))
expected_theta = 0.1 + 0.8 * expected_occ
- assert df["median"].iloc[0] == pytest.approx(expected_theta, rel=1e-5)
+ assert df["q0.5"].iloc[0] == pytest.approx(expected_theta, rel=1e-5)
def test_different_indices_use_different_posteriors(self, model):
"""
@@ -400,11 +400,11 @@ def test_different_indices_use_different_posteriors(self, model):
# iptg(0)/A(1): theta_low=1.0, conc=1.0→log=0, log_K=0, n=1
occ_iptg = 1.0 / (1.0 + np.exp(-1.0 * (np.log(1.0) - 0.0)))
- assert row_iptg_A["median"].iloc[0] == pytest.approx(1.0 + 1.0 * occ_iptg, rel=1e-5)
+ assert row_iptg_A["q0.5"].iloc[0] == pytest.approx(1.0 + 1.0 * occ_iptg, rel=1e-5)
# atc(1)/B(2): theta_low=5.0, conc=2.0→log=log(2), log_K=0, n=1
occ_atc = 1.0 / (1.0 + np.exp(-1.0 * (np.log(2.0) - 0.0)))
- assert row_atc_B["median"].iloc[0] == pytest.approx(5.0 + 1.0 * occ_atc, rel=1e-5)
+ assert row_atc_B["q0.5"].iloc[0] == pytest.approx(5.0 + 1.0 * occ_atc, rel=1e-5)
def test_zero_concentration_handled(self, model):
"""Zero concentrations should not produce NaN (uses ZERO_CONC_VALUE)."""
@@ -415,7 +415,7 @@ def test_zero_concentration_handled(self, model):
})
df = extract_theta_curves(model, self._flat_posteriors(),
q_to_get=_Q, manual_titrant_df=manual_df, num_samples=None)
- assert not df["median"].isna().any()
+ assert not df["q0.5"].isna().any()
def test_broadcast_without_genotype(self, model):
manual_df = pd.DataFrame({
@@ -444,7 +444,7 @@ def test_dispatches_to_hill(self):
}
df = extract_theta_curves(model, posteriors, q_to_get=_Q)
assert "map_theta_group" not in df.columns
- assert "median" in df.columns
+ assert "q0.5" in df.columns
def test_dispatches_to_hill_mut(self):
model = _make_model(theta="hill_mut")
@@ -456,7 +456,7 @@ def test_dispatches_to_hill_mut(self):
}
df = extract_theta_curves(model, posteriors, q_to_get=_Q)
assert "genotype_idx" not in df.columns
- assert "median" in df.columns
+ assert "q0.5" in df.columns
def test_raises_for_categorical(self):
# categorical now supports the interface; use 'simple' to test the error path
@@ -519,5 +519,5 @@ def test_hill_and_hill_mut_agree_on_math(self):
}
df_mut = extract_theta_curves(model_mut, posteriors_mut, q_to_get=_Q)
- assert df_hill["median"].iloc[0] == pytest.approx(
- df_mut["median"].iloc[0], rel=1e-5)
+ assert df_hill["q0.5"].iloc[0] == pytest.approx(
+ df_mut["q0.5"].iloc[0], rel=1e-5)
diff --git a/tests/tfscreen/tfmodel/test_extraction_robustness.py b/tests/tfscreen/tfmodel/test_extraction_robustness.py
index 799fa65a..c3592876 100644
--- a/tests/tfscreen/tfmodel/test_extraction_robustness.py
+++ b/tests/tfscreen/tfmodel/test_extraction_robustness.py
@@ -32,8 +32,8 @@ def mock_model():
"titrant_conc": [1.0],
"replicate": ["1"],
"condition_rep": ["cond1"],
- "condition_pre": ["pre1"],
- "condition_sel": ["sel1"],
+ "condition_pre": ["pre-1"],
+ "condition_sel": ["sel+1"],
"map_theta": [0],
"map_theta_group": [0],
"map_condition_rep": [0],
@@ -50,7 +50,7 @@ def mock_model():
'condition_rep': pd.DataFrame({"replicate": ["1"], "condition_rep": ["cond1"], "map_condition_rep": [0]})
}
mock_tm.tensor_dim_names = ["replicate", "time", "condition_pre", "condition_sel", "titrant_name", "titrant_conc", "genotype"]
- mock_tm.tensor_dim_labels = [["1"], ["1"], ["pre1"], ["sel1"], ["iptg"], [1.0], ["wt"]]
+ mock_tm.tensor_dim_labels = [["1"], ["1"], ["pre-1"], ["sel+1"], ["iptg"], [1.0], ["wt"]]
model.growth_tm = mock_tm
model.training_tm = mock_tm
model.mut_labels = []
@@ -176,7 +176,7 @@ def test_extract_theta_curves_manual_genotype(mock_model):
"theta_theta_low": np.random.rand(10, 1)
}
df = extract_theta_curves(mock_model, posteriors, manual_titrant_df=manual_df)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
assert len(df) == 1
def test_extract_theta_curves_broadcast(mock_model):
@@ -274,9 +274,10 @@ def test_extract_parameters_growth_transition_shares_replicates(
def test_extract_parameters_errors(mock_model):
- """Test error handling in extract_parameters."""
- with pytest.raises(ValueError, match="q_to_get should be a dictionary"):
- extract_parameters(mock_model, {}, q_to_get=[0.5])
+ """Test error handling in extract_parameters: dict is now invalid, list is valid."""
+ # A plain list is now the correct API; a dict should raise
+ with pytest.raises(ValueError):
+ extract_parameters(mock_model, {}, q_to_get={"point_est": 0.5})
def test_extract_parameters_h5_file(mock_model, tmp_path):
"""Test extract_parameters loading from an HDF5 file path."""
@@ -310,7 +311,7 @@ def test_extract_theta_curves_hdf5_path(mock_model, tmp_path):
f.create_dataset(k, data=v)
df = extract_theta_curves(mock_model, h5_path)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
def test_extract_growth_predictions_hdf5_basic(mock_model, tmp_path):
"""Test extract_growth_predictions reads correctly from HDF5."""
@@ -325,11 +326,11 @@ def test_extract_growth_predictions_hdf5_basic(mock_model, tmp_path):
# Test reading from an open h5py.File
with h5py.File(h5_path, "r") as f:
df = extract_growth_predictions(mock_model, f)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
# Test load from file path (.h5) to hit the File loading branch
df_path = extract_growth_predictions(mock_model, h5_path)
- assert "median" in df_path.columns
+ assert "q0.5" in df_path.columns
def test_extract_growth_predictions_hdf5_multiple_groups(mock_model, tmp_path):
"""Test extract_growth_predictions with multiple (rep, time, condition) groups.
@@ -341,8 +342,8 @@ def test_extract_growth_predictions_hdf5_multiple_groups(mock_model, tmp_path):
mock_model.growth_df = pd.DataFrame({
"replicate": ["1", "2"],
"genotype": ["wt", "wt"],
- "condition_pre": ["pre1", "pre2"],
- "condition_sel": ["sel1", "sel2"],
+ "condition_pre": ["pre-1", "pre-2"],
+ "condition_sel": ["sel+1", "sel+2"],
"titrant_name": ["iptg", "iptg"],
"titrant_conc": [1.0, 2.0],
"t_pre": [0.0, 0.0],
@@ -368,14 +369,14 @@ def test_extract_growth_predictions_hdf5_multiple_groups(mock_model, tmp_path):
with h5py.File(h5_path, "r") as f:
df = extract_growth_predictions(mock_model, f)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
assert len(df) == 2
def test_extract_theta_curves_q_to_get_error(mock_model):
- """Test extract_theta_curves with invalid q_to_get."""
+ """Test extract_theta_curves with invalid q_to_get: dict is now invalid."""
mock_model._theta = "hill_geno"
- with pytest.raises(ValueError, match="q_to_get should be a dictionary"):
- extract_theta_curves(mock_model, {}, q_to_get=[0.5])
+ with pytest.raises(ValueError):
+ extract_theta_curves(mock_model, {}, q_to_get={"median": 0.5})
def test_extract_growth_predictions_hdf5_is_h5_check(mock_model, tmp_path):
"""Explicitly check that we are hitting the is_h5 branch."""
@@ -390,4 +391,4 @@ def test_extract_growth_predictions_hdf5_is_h5_check(mock_model, tmp_path):
ds = f["growth_pred"]
# extract_growth_predictions takes param_posteriors which is the file/group
df = extract_growth_predictions(mock_model, f)
- assert "median" in df.columns
+ assert "q0.5" in df.columns
diff --git a/tests/tfscreen/tfmodel/test_fit_model_helpers.py b/tests/tfscreen/tfmodel/test_fit_model_helpers.py
index 65c5e905..cb39838b 100644
--- a/tests/tfscreen/tfmodel/test_fit_model_helpers.py
+++ b/tests/tfscreen/tfmodel/test_fit_model_helpers.py
@@ -43,19 +43,15 @@ def test_run_svi_flow_converged(mock_run_inference):
"""Test standard SVI flow where convergence is reached."""
_, ri = mock_run_inference
- # Execute
- state, params, converged = _run_svi(
+ svi_obj, state, params, converged = _run_svi(
ri,
init_params=None,
out_prefix="test_root",
max_num_epochs=500,
- num_posterior_samples=100
)
- # 1. Setup SVI
ri.setup_svi.assert_called_once()
- # 2. Run Optimization
ri.run_optimization.assert_called_once_with(
"mock_svi_obj",
init_params=None,
@@ -71,45 +67,24 @@ def test_run_svi_flow_converged(mock_run_inference):
epoch_checkpoint_interval=ANY
)
- # 3. Get Posteriors (because converged=True)
- ri.get_posteriors.assert_called_once_with(
- svi="mock_svi_obj",
- svi_state="final_state",
- out_prefix="test_root",
- num_posterior_samples=100,
- sampling_batch_size=ANY,
- forward_batch_size=ANY
- )
+ # Posteriors are never sampled by _run_svi
+ ri.get_posteriors.assert_not_called()
+ # svi_obj is the object returned by ri.setup_svi
+ assert svi_obj == "mock_svi_obj"
assert state == "final_state"
assert converged is True
-def test_run_svi_flow_not_converged_no_posterior(mock_run_inference):
- """Test SVI flow where it fails to converge and always_get_posterior is False."""
+def test_run_svi_flow_not_converged(mock_run_inference):
+ """Test SVI flow where it fails to converge."""
_, ri = mock_run_inference
- # Force non-convergence
ri.run_optimization.return_value = ("state", {}, False)
- state, params, converged = _run_svi(
- ri,
- init_params=None,
- always_get_posterior=False
- )
+ _, state, params, converged = _run_svi(ri, init_params=None)
- # Should NOT get posteriors
ri.get_posteriors.assert_not_called()
assert converged is False
-def test_run_svi_flow_always_posterior(mock_run_inference):
- """Test SVI flow where it fails to converge but posteriors are forced."""
- _, ri = mock_run_inference
- ri.run_optimization.return_value = ("state", {}, False)
-
- _run_svi(ri, init_params=None, always_get_posterior=True)
-
- # Should get posteriors despite no convergence
- ri.get_posteriors.assert_called_once()
-
def test_run_svi_not_converged_stdout(mock_run_inference, capsys):
"""Test SVI not converged message."""
_, ri = mock_run_inference
@@ -131,7 +106,6 @@ def test_run_map_flow(mock_run_inference):
init_params=init_params,
out_prefix="test_map",
max_num_epochs=1000,
- num_posterior_samples=100
)
# 1. Setup MAP
@@ -153,17 +127,16 @@ def test_run_map_flow(mock_run_inference):
epoch_checkpoint_interval=ANY
)
- # 3. Write Params
ri.write_params.assert_called_once_with({"p": 1}, out_prefix="test_map")
- # always_get_posterior=False is default
+ # Posteriors are never sampled by _run_map
ri.get_posteriors.assert_not_called()
def test_run_map_not_converged(mock_run_inference, capsys):
"""Test MAP not converged message."""
_, ri = mock_run_inference
ri.run_optimization.return_value = ("state", {"p": 1}, False)
- _run_map(ri, init_params={"p": 1}, always_get_posterior=False)
+ _run_map(ri, init_params={"p": 1})
captured = capsys.readouterr()
assert "MAP run converged" not in captured.out
assert "MAP run has not yet converged" in captured.out
diff --git a/tests/tfscreen/tfmodel/test_get_posteriors_mapping.py b/tests/tfscreen/tfmodel/test_get_posteriors_mapping.py
index 4ee486d8..37a4aa5d 100644
--- a/tests/tfscreen/tfmodel/test_get_posteriors_mapping.py
+++ b/tests/tfscreen/tfmodel/test_get_posteriors_mapping.py
@@ -123,6 +123,53 @@ def test_get_posteriors_batching_mapping(tmpdir):
np.testing.assert_allclose(det_p_in, post_batched["det_p_in"][:])
np.testing.assert_allclose(det_p_out, post_batched["det_p_out"][:])
+def test_dim_map_excludes_non_genotype_plate_at_same_dim():
+ """A plate at dim=-1 that is NOT the genotype plate must not appear in dim_map.
+
+ Reproduces a bug where hierarchical_factored's tube_offset site (in the
+ condition_pre plate at dim=-1) was falsely identified as genotype-indexed
+ when num_condition_pre == probe_size == 2. The shape mismatch caused:
+ (1, 1, G, 1, 1, 1, G) vs (1, T, CP, 1, 1, TC, G) during posterior sampling.
+ """
+
+ class MockModelWithCondPre:
+ def __init__(self, num_genotype=10, num_cond_pre=2):
+ self.data = MockData(num_genotype=num_genotype,
+ batch_size=num_genotype,
+ batch_idx=jnp.arange(num_genotype))
+ self._num_cond_pre = num_cond_pre
+ self.priors = {}
+
+ def jax_model(self, data, priors):
+ with numpyro.plate("shared_genotype_plate", data.num_genotype, dim=-1):
+ numpyro.sample("geno_p", dist.Normal(0, 1))
+ # Separate plate at the same dim as genotype — must NOT appear in dim_map
+ with numpyro.plate("replicate_plate", 1, dim=-2):
+ with numpyro.plate("cond_pre_plate", self._num_cond_pre, dim=-1):
+ numpyro.sample("tube_offset", dist.Normal(0, 1))
+
+ def jax_model_guide(self, data, priors):
+ return AutoDelta(self.jax_model)(data, priors)
+
+ def get_random_idx(self, key=None, num_batches=1):
+ if num_batches == 1:
+ return np.array([0])
+ return np.zeros((num_batches, 1), dtype=int)
+
+ def get_batch(self, data, indices):
+ return MockData(num_genotype=len(indices),
+ batch_size=len(indices),
+ batch_idx=indices)
+
+ # Use num_genotype=10, num_cond_pre=2 — probe_size becomes min(2,10)=2 == num_cond_pre
+ model = MockModelWithCondPre(num_genotype=10, num_cond_pre=2)
+ ri = RunInference(model, seed=0)
+ dim_map = ri._get_genotype_dim_map()
+
+ assert "geno_p" in dim_map, "genotype-plate site must appear in dim_map"
+ assert "tube_offset" not in dim_map, "non-genotype plate site must not appear in dim_map"
+
+
def test_get_posteriors_shape_ambiguity(tmpdir):
# Test with num_genotypes == num_titrants to ensure it doesn't get "mixed up"
num_genotypes = 3
diff --git a/tests/tfscreen/tfmodel/test_h5_indexing.py b/tests/tfscreen/tfmodel/test_h5_indexing.py
index d2dd67ca..af4c540e 100644
--- a/tests/tfscreen/tfmodel/test_h5_indexing.py
+++ b/tests/tfscreen/tfmodel/test_h5_indexing.py
@@ -14,8 +14,8 @@ def test_predict_h5_indexing(tmp_path):
"genotype": ["wt"],
"titrant_name": ["tit1"],
"titrant_conc": [0.0],
- "condition_pre": ["pre1"],
- "condition_sel": ["sel1"],
+ "condition_pre": ["pre-1"],
+ "condition_sel": ["sel+1"],
"t_pre": [10.0],
"t_sel": [0.0],
"ln_cfu": [0.0],
@@ -29,7 +29,7 @@ def test_predict_h5_indexing(tmp_path):
"theta_obs": [0.5],
"theta_std": [0.01]
})
- mc = ModelOrchestrator(growth_df, binding_df)
+ orchestrator = ModelOrchestrator(growth_df, binding_df)
# 2. Create an H5 file with some "posteriors"
h5_path = tmp_path / "posteriors.h5"
@@ -61,7 +61,7 @@ def test_predict_h5_indexing(tmp_path):
# This calls val = val[sample_indices] inside prediction.py
try:
- results = predict(mc, str(h5_path), num_samples=10)
+ results = predict(orchestrator, str(h5_path), num_samples=10)
except TypeError as e:
if "Indexing elements must be in increasing order" in str(e):
pytest.fail("H5 indexing failed: elements not in increasing order")
diff --git a/tests/tfscreen/tfmodel/test_model_orchestrator.py b/tests/tfscreen/tfmodel/test_model_orchestrator.py
index 9844bf50..fe2a541f 100644
--- a/tests/tfscreen/tfmodel/test_model_orchestrator.py
+++ b/tests/tfscreen/tfmodel/test_model_orchestrator.py
@@ -117,11 +117,30 @@ def test_read_growth_df_no_replicate(mocker, base_growth_df):
mocker.patch("tfscreen.util.dataframe.get_scaled_cfu", return_value=df_no_rep)
mocker.patch("tfscreen.util.dataframe.check_columns")
mocker.patch("tfscreen.tfmodel.model_orchestrator.add_group_columns", return_value=df_no_rep)
-
+
res = _read_growth_df("path.csv")
assert "replicate" in res.columns
assert res["replicate"].iloc[0] == 1
+def test_read_growth_df_coerces_time_to_float(mocker, base_growth_df):
+ """t_pre and t_sel stored as int64 (e.g. from pd.read_csv with integer values)
+ must be cast to float64 so that downstream merges against float prediction grids
+ don't produce int/float dtype warnings."""
+ df_int_time = base_growth_df.copy()
+ df_int_time["t_pre"] = df_int_time["t_pre"].astype(int)
+ df_int_time["t_sel"] = df_int_time["t_sel"].astype(int)
+ assert df_int_time["t_pre"].dtype == int
+ assert df_int_time["t_sel"].dtype == int
+
+ mocker.patch("tfscreen.util.io.read_dataframe", return_value=df_int_time)
+ mocker.patch("tfscreen.util.dataframe.get_scaled_cfu", return_value=df_int_time)
+ mocker.patch("tfscreen.util.dataframe.check_columns")
+ mocker.patch("tfscreen.tfmodel.model_orchestrator.add_group_columns", return_value=df_int_time)
+
+ res = _read_growth_df("path.csv")
+ assert res["t_pre"].dtype == float, f"expected float, got {res['t_pre'].dtype}"
+ assert res["t_sel"].dtype == float, f"expected float, got {res['t_sel'].dtype}"
+
# ----------------------------------------------------------------------------
# 5. Tests for _read_binding_df
# ----------------------------------------------------------------------------
@@ -143,25 +162,60 @@ def test_read_binding_df(mocker):
res = _read_binding_df("path.csv", growth_df=growth_df)
assert "theta_obs" in res.columns
-def test_read_binding_df_errors(mocker):
+def test_read_binding_df_missing_col_error(mocker):
growth_df = pd.DataFrame({"genotype": ["A"], "titrant_name": ["T1"]})
binding_df = pd.DataFrame({
- "genotype": ["B"], "titrant_name": ["T1"],
+ "genotype": ["A"], "titrant_name": ["T1"],
"theta_obs": [0.5], "theta_std": [0.1], "titrant_conc": [1.0]
})
mocker.patch("tfscreen.util.io.read_dataframe", return_value=binding_df)
mocker.patch("tfscreen.genetics.set_categorical_genotype", return_value=binding_df)
- mocker.patch("tfscreen.util.dataframe.check_columns")
-
- with patch("pandas.Index.isin", return_value=np.array([False])):
- with pytest.raises(ValueError, match="not seen"):
- _read_binding_df("path.csv", growth_df=growth_df)
-
mocker.patch("tfscreen.util.dataframe.check_columns", side_effect=ValueError("missing col"))
with pytest.raises(ValueError, match="missing col"):
_read_binding_df("path.csv", growth_df=growth_df)
+def test_read_binding_df_extra_pairs_dropped_with_warning(capsys):
+ """binding rows whose (genotype, titrant_name) is absent from growth_df are dropped."""
+ from tfscreen.genetics import set_categorical_genotype
+ import tfscreen.util.dataframe
+
+ growth_df = pd.DataFrame({
+ "genotype": ["wt", "wt"],
+ "titrant_name": ["iptg", "iptg"],
+ "titrant_conc": [0.0, 1.0],
+ "ln_cfu": [10.0, 10.0],
+ "ln_cfu_std": [0.1, 0.1],
+ "t_pre": [1.0, 1.0],
+ "t_sel": [1.0, 1.0],
+ "replicate": [1, 1],
+ "condition_pre": ["kan", "kan"],
+ "condition_sel": ["kan", "kan"],
+ })
+ growth_df = set_categorical_genotype(growth_df, standardize=True)
+ growth_df = tfscreen.util.dataframe.add_group_columns(
+ growth_df, ["genotype", "titrant_name"], "map_theta_group"
+ )
+
+ # binding has "wt" (in growth) and "M42I" (not in growth)
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "M42I"],
+ "titrant_name": ["iptg", "iptg"],
+ "titrant_conc": [1.0, 1.0],
+ "theta_obs": [0.1, 0.5],
+ "theta_std": [0.02, 0.02],
+ })
+
+ result = _read_binding_df(binding_df, growth_df=growth_df)
+ captured = capsys.readouterr()
+ assert "will be dropped" in captured.out
+ assert "M42I" in captured.out
+
+ # Only the "wt" row should survive
+ assert set(result["genotype"]) == {"wt"}
+ assert len(result) == 1
+
+
def test_read_binding_df_preserves_canonical_genotype_order():
# Regression test: the pd.merge inside add_group_columns converts the
# categorical genotype column to object dtype. _read_binding_df must
@@ -453,8 +507,8 @@ def test_extract_parameters_full(initialized_model_class):
def test_extract_parameters_errors(initialized_model_class):
model = initialized_model_class
- with pytest.raises(ValueError, match="should be a dictionary"):
- extract_parameters(model, {}, q_to_get="not_a_dict")
+ with pytest.raises(ValueError, match="should be a 1-D array-like"):
+ extract_parameters(model, {}, q_to_get={"not": "valid"})
def test_extract_parameters_npz(initialized_model_class, tmpdir):
model = initialized_model_class
@@ -530,8 +584,8 @@ def test_extract_theta_curves_errors(initialized_model_class):
with pytest.raises(ValueError, match="does not support this interface"):
extract_theta_curves(model, {})
model._theta = "hill_geno"
- with pytest.raises(ValueError, match="should be a dictionary"):
- extract_theta_curves(model, {}, q_to_get="not_a_dict")
+ with pytest.raises(ValueError, match="should be a 1-D array-like"):
+ extract_theta_curves(model, {}, q_to_get={"not": "valid"})
def test_extract_growth_predictions_full(initialized_model_class, tmpdir):
model = initialized_model_class
@@ -557,12 +611,12 @@ def test_extract_growth_predictions_full(initialized_model_class, tmpdir):
})
post = {"growth_pred": np.zeros((1, 1, 1, 1, 1, 1, 1, 1))}
res = extract_growth_predictions(model, post)
- assert "median" in res.columns
-
+ assert "q0.5" in res.columns
+
with pytest.raises(ValueError, match="'growth_pred' not found"):
extract_growth_predictions(model, {})
- with pytest.raises(ValueError, match="should be a dictionary"):
- extract_growth_predictions(model, {"growth_pred": None}, q_to_get="bad")
+ with pytest.raises(ValueError, match="should be a 1-D array-like"):
+ extract_growth_predictions(model, {"growth_pred": None}, q_to_get={"not": "valid"})
path = os.path.join(tmpdir, "gp.npz")
np.savez(path, growth_pred=np.zeros((1,1,1,1,1,1,1,1)))
@@ -762,4 +816,154 @@ def capture_populate(cls, sources):
assert "scale_vector" in captured, "BindingData populate_dataclass was never called"
expected = base_scale * explicit_weight
- np.testing.assert_allclose(captured["scale_vector"], expected)
\ No newline at end of file
+ np.testing.assert_allclose(captured["scale_vector"], expected)
+
+
+# ---------------------------------------------------------------------------
+# Tests for condition-label extraction logic in _setup_model
+# ---------------------------------------------------------------------------
+
+class TestConditionLabelExtraction:
+ """
+ Verify that _setup_model correctly extracts ordered condition labels from
+ growth_tm and routes them to get_priors().
+
+ These tests exercise the extraction logic directly without needing to
+ build a full ModelOrchestrator, by calling the relevant code path via
+ a stripped-down mock orchestrator.
+ """
+
+ def _make_cond_rep_df(self, names, indices):
+ """Return a condition_rep map_groups DataFrame."""
+ return pd.DataFrame({
+ "condition_rep": names,
+ "map_condition_rep": indices,
+ })
+
+ def test_labels_extracted_in_index_order(self):
+ """
+ Labels must be sorted by map_condition_rep so parameter positions
+ match the ordering established by the TensorManager.
+ """
+ # Rows are intentionally NOT in index order
+ df = self._make_cond_rep_df(
+ names=["pheS+4CP", "pheS-4CP", "kanR+kan", "kanR-kan"],
+ indices=[2, 3, 0, 1],
+ )
+ result = list(df.sort_values("map_condition_rep")["condition_rep"])
+ assert result == ["kanR+kan", "kanR-kan", "pheS+4CP", "pheS-4CP"]
+
+ def test_labels_correct_when_already_ordered(self):
+ df = self._make_cond_rep_df(
+ names=["kanR+kan", "kanR-kan", "pheS+4CP", "pheS-4CP"],
+ indices=[0, 1, 2, 3],
+ )
+ result = list(df.sort_values("map_condition_rep")["condition_rep"])
+ assert result == ["kanR+kan", "kanR-kan", "pheS+4CP", "pheS-4CP"]
+
+ def test_single_condition(self):
+ df = self._make_cond_rep_df(["sel+media"], [0])
+ result = list(df.sort_values("map_condition_rep")["condition_rep"])
+ assert result == ["sel+media"]
+
+ def test_shares_replicates_false_includes_replicate_but_still_sorts(self):
+ """When shares_replicates=False the df has (replicate, condition_rep) rows."""
+ df = pd.DataFrame({
+ "replicate": [1, 2, 1, 2],
+ "condition_rep": ["pheS+4CP", "pheS+4CP", "pheS-4CP", "pheS-4CP"],
+ "map_condition_rep": [1, 3, 0, 2],
+ })
+ result = list(df.sort_values("map_condition_rep")["condition_rep"])
+ # Sorted by index: 0=pheS-4CP, 1=pheS+4CP, 2=pheS-4CP, 3=pheS+4CP
+ assert result[0] == "pheS-4CP"
+ assert result[1] == "pheS+4CP"
+
+ def test_linear_get_priors_signature_has_condition_labels(self):
+ """
+ The orchestrator detects condition_labels support via inspect.signature.
+ Verify that the linear component's get_priors actually has this param.
+ """
+ import inspect
+ from tfscreen.tfmodel.generative.components.growth.linear import get_priors
+ sig = inspect.signature(get_priors)
+ assert "condition_labels" in sig.parameters, (
+ "linear.get_priors must have a 'condition_labels' parameter so "
+ "ModelOrchestrator can pass condition names for per-condition m priors"
+ )
+
+ def test_setup_model_extracts_labels_and_calls_get_priors(self, mocker):
+ """
+ _setup_model must call get_priors(condition_labels=...) for a
+ condition_growth component that accepts that parameter.
+ """
+ import inspect
+ import tfscreen.tfmodel.model_orchestrator as mo
+ from tfscreen.tfmodel.generative.components.growth import linear as real_linear
+
+ captured = {}
+
+ def capturing_get_priors(condition_labels=None):
+ captured["condition_labels"] = condition_labels
+ return real_linear.get_priors(condition_labels=condition_labels)
+
+ cond_rep_df = pd.DataFrame({
+ "condition_rep": ["kanR+kan", "kanR-kan"],
+ "map_condition_rep": [0, 1],
+ })
+ mock_growth_tm = MagicMock()
+ mock_growth_tm.map_groups = {"condition_rep": cond_rep_df}
+
+ # Stub ALL registry entries except we intercept condition_growth.get_priors
+ all_stubs = {}
+ for comp_key, comp_val in mo.model_registry.items():
+ if not isinstance(comp_val, dict):
+ continue
+ all_stubs[comp_key] = {}
+ for variant_name, variant_mod in comp_val.items():
+ stub = MagicMock()
+ stub.get_priors.return_value = MagicMock()
+ stub.get_guesses.return_value = {}
+ stub.define_model = MagicMock()
+ stub.guide = MagicMock()
+ stub.calculate_growth = MagicMock()
+ stub.rescale = MagicMock()
+ stub.observe = MagicMock()
+ all_stubs[comp_key][variant_name] = stub
+
+ # Override condition_growth / linear with our spy
+ all_stubs.setdefault("condition_growth", {})["linear"] = MagicMock(
+ get_priors=capturing_get_priors,
+ get_guesses=MagicMock(return_value={}),
+ define_model=MagicMock(),
+ guide=MagicMock(),
+ calculate_growth=MagicMock(),
+ )
+
+ orchestrator = MagicMock(spec=ModelOrchestrator)
+ orchestrator.growth_tm = mock_growth_tm
+ orchestrator._data = MagicMock()
+ orchestrator._binding_only = False
+ orchestrator._batch_size = None
+ orchestrator._condition_growth = "linear"
+ orchestrator._growth_transition = "instant"
+ orchestrator._ln_cfu0 = "hierarchical"
+ orchestrator._dk_geno = "fixed"
+ orchestrator._activity = "fixed"
+ orchestrator._theta = "hill_geno"
+ orchestrator._transformation = "single"
+ orchestrator._theta_growth_noise = "zero"
+ orchestrator._theta_binding_noise = "zero"
+ orchestrator._growth_noise = "zero"
+ orchestrator._sample_offset = "zero"
+ orchestrator._theta_rescale = "passthrough"
+
+ with patch.object(mo, "model_registry", all_stubs):
+ try:
+ mo.ModelOrchestrator._initialize_classes(orchestrator)
+ except Exception:
+ pass # only condition_labels matters
+
+ assert "condition_labels" in captured, (
+ "_setup_model did not pass condition_labels to get_priors"
+ )
+ assert captured["condition_labels"] == ["kanR+kan", "kanR-kan"]
\ No newline at end of file
diff --git a/tests/tfscreen/tfmodel/test_needs_mut_wiring.py b/tests/tfscreen/tfmodel/test_needs_mut_wiring.py
index 682f9ba1..3a23842c 100644
--- a/tests/tfscreen/tfmodel/test_needs_mut_wiring.py
+++ b/tests/tfscreen/tfmodel/test_needs_mut_wiring.py
@@ -52,9 +52,9 @@ def _make_growth_csv(tmp_path, genotypes=_GENOTYPES):
"num_muts": num_muts,
"counts": 1000,
"replicate": rep,
- "condition_pre": "pre",
+ "condition_pre": "pre-cond",
"t_pre": 24.0,
- "condition_sel": "sel",
+ "condition_sel": "sel+cond",
"t_sel": 48.0,
"titrant_name": "iptg",
"titrant_conc": conc,
@@ -115,11 +115,11 @@ def test_needs_mut_growth_path(tmp_path, theta):
growth_path = _make_growth_csv(tmp_path)
binding_path = _make_binding_csv(tmp_path)
- mc = ModelOrchestrator(growth_path, binding_path, theta=theta)
+ orchestrator = ModelOrchestrator(growth_path, binding_path, theta=theta)
- assert len(mc.mut_labels) > 0, \
+ assert len(orchestrator.mut_labels) > 0, \
f"{theta} (growth): mut_labels is empty — _needs_mut did not fire"
- assert mc.data.growth.num_mutation > 0, \
+ assert orchestrator.data.growth.num_mutation > 0, \
f"{theta} (growth): num_mutation == 0 — mutation matrices not built"
@@ -141,11 +141,11 @@ def test_needs_mut_binding_only_path(tmp_path, theta):
binding_path = _make_binding_csv(tmp_path)
- mc = ModelOrchestrator(None, binding_path, binding_only=True, theta=theta)
+ orchestrator = ModelOrchestrator(None, binding_path, binding_only=True, theta=theta)
- assert len(mc.mut_labels) > 0, \
+ assert len(orchestrator.mut_labels) > 0, \
f"{theta} (binding_only): mut_labels is empty — _needs_mut did not fire"
- assert mc.data.binding.num_mutation > 0, \
+ assert orchestrator.data.binding.num_mutation > 0, \
f"{theta} (binding_only): num_mutation == 0 — mutation matrices not built"
@@ -165,13 +165,13 @@ def test_offset_guesses_have_correct_shape_binding_only(tmp_path, theta):
from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
binding_path = _make_binding_csv(tmp_path)
- mc = ModelOrchestrator(None, binding_path, binding_only=True, theta=theta)
+ orchestrator = ModelOrchestrator(None, binding_path, binding_only=True, theta=theta)
- M = mc.data.binding.num_mutation
+ M = orchestrator.data.binding.num_mutation
assert M > 0
# Every *_offset key in init_params must have length M (or M as a dimension)
- for key, val in mc.init_params.items():
+ for key, val in orchestrator.init_params.items():
if "offset" in key and hasattr(val, "shape") and val.shape != ():
assert val.shape[0] == M or val.shape[-1] == M, \
f"{key} has shape {val.shape} but expected M={M} on at least one axis"
diff --git a/tests/tfscreen/tfmodel/test_posteriors.py b/tests/tfscreen/tfmodel/test_posteriors.py
index 6ac90a49..127e87f2 100644
--- a/tests/tfscreen/tfmodel/test_posteriors.py
+++ b/tests/tfscreen/tfmodel/test_posteriors.py
@@ -7,11 +7,11 @@ def test_load_posteriors_dict():
"""Test loading posteriors from a dictionary."""
posteriors = {"param1": np.array([1, 2, 3]), "param2": np.array([4, 5, 6])}
q, p = load_posteriors(posteriors)
-
+
assert p == posteriors
- assert "median" in q
- assert q["median"] == 0.5
- assert len(q) == 9
+ assert "q0.5" in q
+ assert q["q0.5"] == 0.5
+ assert len(q) == 17
def test_load_posteriors_npz(tmp_path):
"""Test loading posteriors from an .npz file."""
@@ -22,7 +22,7 @@ def test_load_posteriors_npz(tmp_path):
assert "param1" in p
assert np.array_equal(p["param1"], [1, 2, 3])
- assert "median" in q
+ assert "q0.5" in q
def test_load_posteriors_h5(tmp_path):
"""Test loading posteriors from an .h5 file."""
@@ -45,19 +45,19 @@ def test_load_posteriors_h5(tmp_path):
def test_load_posteriors_custom_q():
"""Test loading with custom quantiles."""
posteriors = {"param1": np.array([1, 2, 3])}
- custom_q = {"test": 0.123}
- q, p = load_posteriors(posteriors, q_to_get=custom_q)
-
- assert q == custom_q
+ q, p = load_posteriors(posteriors, q_to_get=[0.123])
+
+ assert "q0.123" in q
+ assert q["q0.123"] == 0.123
assert p == posteriors
def test_load_posteriors_errors():
"""Test validation errors."""
posteriors = {"param1": np.array([1, 2, 3])}
-
- with pytest.raises(ValueError, match="q_to_get should be a dictionary"):
- load_posteriors(posteriors, q_to_get=[0.5])
-
+
+ with pytest.raises(ValueError, match="q_to_get should be a 1-D array-like"):
+ load_posteriors(posteriors, q_to_get={"test": 0.5})
+
# Test invalid file path (np.load will raise error)
with pytest.raises(Exception):
load_posteriors("non_existent_file.npz")
diff --git a/tests/tfscreen/tfmodel/test_prediction.py b/tests/tfscreen/tfmodel/test_prediction.py
index f66e8ff8..05ce9041 100644
--- a/tests/tfscreen/tfmodel/test_prediction.py
+++ b/tests/tfscreen/tfmodel/test_prediction.py
@@ -1,18 +1,22 @@
import pytest
import pandas as pd
import numpy as np
+import jax.numpy as jnp
+import numpyro.distributions as dist
+from unittest.mock import MagicMock, patch
+
from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
-from tfscreen.tfmodel.analysis.prediction import copy_model_class
+from tfscreen.tfmodel.analysis.prediction import copy_orchestrator, _convert_map_params, _align_site_to_tm_dims, predict
@pytest.fixture
-def dummy_mc():
+def dummy_orchestrator():
"""Create a dummy ModelOrchestrator instance for testing."""
growth_df = pd.DataFrame({
"genotype": ["wt", "wt", "M42V", "M42V"],
"titrant_name": ["tit1", "tit1", "tit1", "tit1"],
"titrant_conc": [0.0, 1.0, 0.0, 1.0],
- "condition_pre": ["pre1", "pre1", "pre1", "pre1"],
- "condition_sel": ["sel1", "sel1", "sel1", "sel1"],
+ "condition_pre": ["pre-1", "pre-1", "pre-1", "pre-1"],
+ "condition_sel": ["sel+1", "sel+1", "sel+1", "sel+1"],
"t_pre": [10.0, 10.0, 10.0, 10.0],
"t_sel": [0.0, 20.0, 0.0, 20.0],
"ln_cfu": [0.0, 5.0, 0.0, 3.0],
@@ -30,64 +34,462 @@ def dummy_mc():
return ModelOrchestrator(growth_df, binding_df)
-def test_copy_model_class_defaults(dummy_mc):
- """Test copy_model_class with all None inputs."""
- new_mc = copy_model_class(dummy_mc)
+def test_copy_orchestrator_defaults(dummy_orchestrator):
+ """Test copy_orchestrator with all None inputs."""
+ new_orchestrator = copy_orchestrator(dummy_orchestrator)
# Categorical combinations (1 rep * 1 cp * 1 cs * 1 tn * 2 geno) = 2
# Quantitative (1 t_pre * 2 t_sel * 2 conc) = 4
# Total = 2 * 4 = 8
- assert len(new_mc.growth_df) == 8
- assert new_mc.settings == dummy_mc.settings
+ assert len(new_orchestrator.growth_df) == 8
+ assert new_orchestrator.settings == dummy_orchestrator.settings
-def test_copy_model_class_expansion(dummy_mc):
+def test_copy_orchestrator_expansion(dummy_orchestrator):
"""Test expanding quantitative inputs."""
- new_mc = copy_model_class(
- dummy_mc,
+ new_orchestrator = copy_orchestrator(
+ dummy_orchestrator,
t_sel=[0.0, 10.0, 20.0],
titrant_conc=[0.0, 0.5, 1.0]
)
# 2 (genotypes) * 3 (t_sel) * 3 (conc) = 18
- assert len(new_mc.growth_df) == 18
- assert set(new_mc.growth_df["t_sel"]) == {0.0, 10.0, 20.0}
- assert set(new_mc.growth_df["titrant_conc"]) == {0.0, 0.5, 1.0}
+ assert len(new_orchestrator.growth_df) == 18
+ assert set(new_orchestrator.growth_df["t_sel"]) == {0.0, 10.0, 20.0}
+ assert set(new_orchestrator.growth_df["titrant_conc"]) == {0.0, 0.5, 1.0}
-def test_copy_model_class_list_inputs(dummy_mc):
+def test_copy_orchestrator_list_inputs(dummy_orchestrator):
"""Test passing single values as non-lists."""
- new_mc = copy_model_class(
- dummy_mc,
+ new_orchestrator = copy_orchestrator(
+ dummy_orchestrator,
t_sel=30.0,
titrant_conc=2.0
)
- assert 30.0 in new_mc.growth_df["t_sel"].values
- assert 2.0 in new_mc.growth_df["titrant_conc"].values
+ assert 30.0 in new_orchestrator.growth_df["t_sel"].values
+ assert 2.0 in new_orchestrator.growth_df["titrant_conc"].values
-def test_copy_model_class_quantitative_errors(dummy_mc):
+def test_copy_orchestrator_quantitative_errors(dummy_orchestrator):
"""Test validation of quantitative inputs."""
with pytest.raises(ValueError, match="t_pre must be >= 0"):
- copy_model_class(dummy_mc, t_pre=-1.0)
+ copy_orchestrator(dummy_orchestrator, t_pre=-1.0)
with pytest.raises(ValueError, match="t_sel must be >= 0"):
- copy_model_class(dummy_mc, t_sel=[-1.0, 1.0])
+ copy_orchestrator(dummy_orchestrator, t_sel=[-1.0, 1.0])
with pytest.raises(ValueError, match="titrant_conc must be >= 0"):
- copy_model_class(dummy_mc, titrant_conc=[1.0, -0.5])
+ copy_orchestrator(dummy_orchestrator, titrant_conc=[1.0, -0.5])
-def test_copy_model_class_t_pre_expansion(dummy_mc):
+def test_copy_orchestrator_t_pre_expansion(dummy_orchestrator):
"""Test expanding t_pre."""
- new_mc = copy_model_class(
- dummy_mc,
+ new_orchestrator = copy_orchestrator(
+ dummy_orchestrator,
t_pre=[10.0, 20.0]
)
# 2 (genotypes) * 2 (t_pre) * 2 (t_sel) * 2 (conc) = 16
- assert len(new_mc.growth_df) == 16
- assert set(new_mc.growth_df["t_pre"]) == {10.0, 20.0}
+ assert len(new_orchestrator.growth_df) == 16
+ assert set(new_orchestrator.growth_df["t_pre"]) == {10.0, 20.0}
-def test_copy_model_class_no_subsetting_binding(dummy_mc):
+def test_copy_orchestrator_no_subsetting_binding(dummy_orchestrator):
"""Test that binding_df is NOT subsetted during copy."""
# Original binding_df has wt and M42V
- assert len(dummy_mc.binding_df) == 2
-
- new_mc = copy_model_class(dummy_mc)
-
+ assert len(dummy_orchestrator.binding_df) == 2
+
+ new_orchestrator = copy_orchestrator(dummy_orchestrator)
+
# New binding_df should still have all genotypes
- assert len(new_mc.binding_df) == 2
+ assert len(new_orchestrator.binding_df) == 2
+
+
+def test_copy_orchestrator_genotypes_subset(dummy_orchestrator):
+ """Requesting a subset of genotypes restricts growth_df rows."""
+ new_orch = copy_orchestrator(dummy_orchestrator, genotypes=["wt"])
+ assert set(new_orch.growth_df["genotype"]) == {"wt"}
+
+
+def test_copy_orchestrator_genotypes_subset_tm_labels(dummy_orchestrator):
+ """TM genotype labels match the requested subset."""
+ new_orch = copy_orchestrator(dummy_orchestrator, genotypes=["wt"])
+ labels = new_orch.growth_tm.tensor_dim_labels[-1].tolist()
+ assert labels == ["wt"]
+
+
+def test_copy_orchestrator_genotypes_subset_row_count(dummy_orchestrator):
+ """Row count matches the single-genotype cross-product."""
+ new_orch = copy_orchestrator(dummy_orchestrator, genotypes=["wt"])
+ # 1 genotype * 1 t_pre * 2 t_sel * 2 titrant_conc = 4
+ assert len(new_orch.growth_df) == 4
+
+
+def test_copy_orchestrator_genotypes_all_explicit(dummy_orchestrator):
+ """Passing all genotypes explicitly is equivalent to genotypes=None."""
+ new_all = copy_orchestrator(dummy_orchestrator)
+ new_explicit = copy_orchestrator(dummy_orchestrator, genotypes=["wt", "M42V"])
+ assert set(new_explicit.growth_df["genotype"]) == set(new_all.growth_df["genotype"])
+
+
+def test_copy_orchestrator_genotypes_unknown_raises(dummy_orchestrator):
+ """Requesting an unknown genotype raises ValueError."""
+ with pytest.raises(ValueError, match="not found"):
+ copy_orchestrator(dummy_orchestrator, genotypes=["no_such_geno"])
+
+
+def test_copy_orchestrator_genotypes_partial_unknown_raises(dummy_orchestrator):
+ """A mix of valid and unknown genotypes still raises."""
+ with pytest.raises(ValueError, match="not found"):
+ copy_orchestrator(dummy_orchestrator, genotypes=["wt", "no_such_geno"])
+
+
+def test_copy_orchestrator_genotypes_tm_labels_differ_from_original(dummy_orchestrator):
+ """Subset TM labels must differ from the original so slicing fires in predict()."""
+ new_orch = copy_orchestrator(dummy_orchestrator, genotypes=["wt"])
+ orig_labels = dummy_orchestrator.growth_tm.tensor_dim_labels[-1].tolist()
+ new_labels = new_orch.growth_tm.tensor_dim_labels[-1].tolist()
+ assert orig_labels != new_labels
+
+
+@pytest.fixture
+def spiked_orchestrator():
+ """ModelOrchestrator with one spiked genotype (A1G) and two library genotypes."""
+ growth_df = pd.DataFrame({
+ "genotype": ["wt", "wt", "M42V", "M42V", "A1G", "A1G"],
+ "titrant_name": ["tit1"] * 6,
+ "titrant_conc": [0.0, 1.0] * 3,
+ "condition_pre": ["pre-1"] * 6,
+ "condition_sel": ["sel+1"] * 6,
+ "t_pre": [10.0] * 6,
+ "t_sel": [0.0, 20.0] * 3,
+ "ln_cfu": [0.0, 5.0, 0.0, 3.0, 0.0, 4.0],
+ "ln_cfu_std": [0.1] * 6,
+ "replicate": [1] * 6,
+ })
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "M42V", "A1G"],
+ "titrant_name": ["tit1"] * 3,
+ "titrant_conc": [0.5] * 3,
+ "theta_obs": [0.5, 0.2, 0.9],
+ "theta_std": [0.01] * 3,
+ })
+ return ModelOrchestrator(growth_df, binding_df, spiked_genotypes=["A1G"])
+
+
+def test_copy_orchestrator_spiked_outside_subset_dropped(spiked_orchestrator):
+ """Spiked genotypes not in the requested subset are removed from settings."""
+ new_orch = copy_orchestrator(spiked_orchestrator, genotypes=["wt", "M42V"])
+ assert new_orch.settings.get("spiked_genotypes") == []
+
+
+def test_copy_orchestrator_spiked_inside_subset_retained(spiked_orchestrator):
+ """Spiked genotypes that are in the requested subset are kept in settings."""
+ new_orch = copy_orchestrator(spiked_orchestrator, genotypes=["wt", "A1G"])
+ assert "A1G" in (new_orch.settings.get("spiked_genotypes") or [])
+
+
+def test_copy_orchestrator_spiked_partial_overlap():
+ """Only spiked genotypes present in the subset survive."""
+ # Build a fresh orchestrator with two spiked genotypes: A1G and S5T
+ growth_df = pd.DataFrame({
+ "genotype": ["wt", "wt", "M42V", "M42V", "A1G", "A1G", "S5T", "S5T"],
+ "titrant_name": ["tit1"] * 8,
+ "titrant_conc": [0.0, 1.0] * 4,
+ "condition_pre": ["pre-1"] * 8,
+ "condition_sel": ["sel+1"] * 8,
+ "t_pre": [10.0] * 8,
+ "t_sel": [0.0, 20.0] * 4,
+ "ln_cfu": [0.0, 5.0, 0.0, 3.0, 0.0, 4.0, 0.0, 4.5],
+ "ln_cfu_std": [0.1] * 8,
+ "replicate": [1] * 8,
+ })
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "M42V", "A1G"],
+ "titrant_name": ["tit1"] * 3,
+ "titrant_conc": [0.5] * 3,
+ "theta_obs": [0.5, 0.2, 0.9],
+ "theta_std": [0.01] * 3,
+ })
+ orch2 = ModelOrchestrator(growth_df, binding_df, spiked_genotypes=["A1G", "S5T"])
+ new_orch = copy_orchestrator(orch2, genotypes=["wt", "A1G"])
+ spiked_kept = new_orch.settings.get("spiked_genotypes") or []
+ assert "A1G" in spiked_kept
+ assert "S5T" not in spiked_kept
+
+
+def test_copy_orchestrator_no_spiked_in_settings_unaffected(dummy_orchestrator):
+ """When spiked_genotypes is absent from settings, subsetting doesn't break."""
+ assert not dummy_orchestrator.settings.get("spiked_genotypes")
+ new_orch = copy_orchestrator(dummy_orchestrator, genotypes=["wt"])
+ assert set(new_orch.growth_df["genotype"]) == {"wt"}
+
+
+def test_copy_orchestrator_spiked_no_subset_unchanged(spiked_orchestrator):
+ """Without genotype subsetting, spiked_genotypes in settings are unchanged."""
+ new_orch = copy_orchestrator(spiked_orchestrator)
+ assert new_orch.settings.get("spiked_genotypes") == \
+ spiked_orchestrator.settings.get("spiked_genotypes")
+
+
+# ---------------------------------------------------------------------------
+# _convert_map_params
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def fake_trace():
+ """
+ Minimal model trace for testing _convert_map_params.
+
+ Includes:
+ - an unconstrained site (Normal → real support)
+ - a positive-constrained site (HalfNormal → positive support)
+ - an observed site (must be passed through with raw value)
+ - a deterministic site (must be ignored — no _auto_loc key)
+ """
+ return {
+ "real_site": {
+ "type": "sample",
+ "is_observed": False,
+ "fn": dist.Normal(0.0, 1.0),
+ },
+ "positive_site": {
+ "type": "sample",
+ "is_observed": False,
+ "fn": dist.HalfNormal(1.0),
+ },
+ "observed_site": {
+ "type": "sample",
+ "is_observed": True,
+ "fn": dist.Normal(0.0, 1.0),
+ },
+ "det_site": {
+ "type": "deterministic",
+ },
+ }
+
+
+class TestConvertMapParams:
+ """Unit tests for _convert_map_params."""
+
+ def test_keys_are_renamed(self, fake_trace):
+ """_auto_loc suffix is stripped from output keys."""
+ map_params = {"real_site_auto_loc": np.array(0.5)}
+ out = _convert_map_params(map_params, fake_trace)
+ assert "real_site" in out
+ assert "real_site_auto_loc" not in out
+
+ def test_leading_sample_dim_added(self, fake_trace):
+ """Each output value has a leading sample dimension of size 1."""
+ map_params = {
+ "real_site_auto_loc": np.array(0.5),
+ "positive_site_auto_loc": np.array(-1.0),
+ }
+ out = _convert_map_params(map_params, fake_trace)
+ for site_name, val in out.items():
+ assert val.shape[0] == 1, (
+ f"{site_name}: expected leading dim 1, got shape {val.shape}"
+ )
+
+ def test_identity_for_unconstrained_site(self, fake_trace):
+ """Normal (real support) site value is passed through unchanged."""
+ raw_val = np.array(-2.3)
+ map_params = {"real_site_auto_loc": raw_val}
+ out = _convert_map_params(map_params, fake_trace)
+ np.testing.assert_allclose(
+ np.asarray(out["real_site"][0]),
+ raw_val,
+ rtol=1e-5,
+ )
+
+ def test_positive_site_passes_through_unchanged(self, fake_trace):
+ """HalfNormal (positive support) site value is passed through as-is.
+
+ RunInference.write_params() calls svi.get_params(), which applies
+ constrain_fn before saving, so stored _auto_loc values are already in
+ constrained (positive) space. _convert_map_params must NOT apply a
+ second bijection.
+
+ Regression: a previous version incorrectly called biject_to(positive)
+ (== ExpTransform) on the stored value, mapping e.g. 2.09 → exp(2.09)
+ = 8.09 instead of leaving it as 2.09. This corrupted ln_cfu0 for
+ library genotypes (which use a sampled HalfNormal scale) while leaving
+ wt predictions unaffected (wt uses a fixed prior scale).
+ """
+ # Use the real-world value from prefit9 that exposed the bug.
+ raw_val = np.array(2.09) # already constrained (positive)
+ map_params = {"positive_site_auto_loc": raw_val}
+ out = _convert_map_params(map_params, fake_trace)
+ constrained = float(out["positive_site"][0])
+
+ # Value must pass through unchanged — no bijection should be applied.
+ np.testing.assert_allclose(constrained, float(raw_val), rtol=1e-5,
+ err_msg="positive site value should not be re-transformed")
+
+ # Explicit guard against the old bug: exp(2.09) ≈ 8.09, not 2.09.
+ assert abs(constrained - float(np.exp(raw_val))) > 1.0, (
+ f"Output ({constrained:.4f}) looks like exp(input) "
+ f"({float(np.exp(raw_val)):.4f}) — biject_to was probably re-introduced"
+ )
+
+ def test_observed_site_ignored(self, fake_trace):
+ """Keys whose site is marked is_observed are not converted."""
+ map_params = {
+ "real_site_auto_loc": np.array(0.0),
+ "observed_site_auto_loc": np.array(1.0),
+ }
+ out = _convert_map_params(map_params, fake_trace)
+ # Both sites should appear with their raw values and a leading sample dim.
+ assert "real_site" in out
+ assert "observed_site" in out
+ assert out["observed_site"].shape[0] == 1
+
+ def test_non_auto_loc_keys_skipped(self, fake_trace):
+ """Keys without _auto_loc suffix are not included in the output."""
+ map_params = {
+ "real_site_auto_loc": np.array(0.0),
+ "some_other_key": np.array(99.0),
+ }
+ out = _convert_map_params(map_params, fake_trace)
+ assert "some_other_key" not in out
+ assert "real_site" in out
+
+ def test_array_site_shape_preserved(self, fake_trace):
+ """Multi-element sites get shape (1, N) not (N,)."""
+ map_params = {"real_site_auto_loc": np.array([1.0, 2.0, 3.0])}
+ out = _convert_map_params(map_params, fake_trace)
+ assert out["real_site"].shape == (1, 3)
+
+ def test_empty_input_returns_empty(self, fake_trace):
+ out = _convert_map_params({}, fake_trace)
+ assert out == {}
+
+ def test_unknown_site_passes_value_through(self, fake_trace):
+ """A key whose site is absent from the trace is still converted (no bijection)."""
+ map_params = {"mystery_site_auto_loc": np.array(5.0)}
+ out = _convert_map_params(map_params, fake_trace)
+ assert "mystery_site" in out
+ np.testing.assert_allclose(float(out["mystery_site"][0]), 5.0)
+
+
+# ---------------------------------------------------------------------------
+# _align_site_to_tm_dims
+# ---------------------------------------------------------------------------
+
+# TM dim order: rep=2, t=3, cp=2, cs=1, tn=1, tc=8, geno=473
+_TM_SIZES = (2, 3, 2, 1, 1, 8, 473)
+
+
+class TestAlignSiteToTmDims:
+ def test_full_rank_identity(self):
+ """A 7-D site matching TM exactly maps to [0,1,2,3,4,5,6]."""
+ spatial = (2, 3, 2, 1, 1, 8, 473)
+ assert _align_site_to_tm_dims(spatial, _TM_SIZES) == [0, 1, 2, 3, 4, 5, 6]
+
+ def test_tail_aligned_1d(self):
+ """A 1-D site of size 473 should map to the last TM dim (geno)."""
+ assert _align_site_to_tm_dims((473,), _TM_SIZES) == [6]
+
+ def test_tail_aligned_2d(self):
+ """(8, 473) should map to TM dims [5, 6] (titrant_conc, geno)."""
+ assert _align_site_to_tm_dims((8, 473), _TM_SIZES) == [5, 6]
+
+ def test_broadcast_left_padding(self):
+ """theta_growth_pred shape (1,1,1,1,1,8,473) — non-broadcast tail match."""
+ spatial = (1, 1, 1, 1, 1, 8, 473)
+ result = _align_site_to_tm_dims(spatial, _TM_SIZES)
+ # Size-1 dims use right-to-left (modulo-safe); non-broadcast dims are [5, 6].
+ assert result[5] == 5
+ assert result[6] == 6
+
+ def test_ln_cfu0_shape(self):
+ """ln_cfu0 shape (2, 2, 473) should map to TM dims [0, 2, 6] = rep, cp, geno."""
+ assert _align_site_to_tm_dims((2, 2, 473), _TM_SIZES) == [0, 2, 6]
+
+ def test_ln_cfu0_tube_offset_shape(self):
+ """(n_rep, n_cp) = (2, 2) maps to TM dims [0, 2]."""
+ assert _align_site_to_tm_dims((2, 2), _TM_SIZES) == [0, 2]
+
+ def test_single_geno_dim(self):
+ """1-D (473,) still maps to geno regardless of other TM sizes."""
+ tm = (5, 3, 7, 473)
+ assert _align_site_to_tm_dims((473,), tm) == [3]
+
+ def test_fallback_to_right_to_left_on_no_match(self):
+ """If subsequence matching fails, right-to-left is returned."""
+ # spatial size 999 is not in any TM dim
+ result = _align_site_to_tm_dims((999,), _TM_SIZES)
+ assert result == [6] # right-to-left: len=7, n_sp=1 → 7-1+0=6
+
+
+# ---------------------------------------------------------------------------
+# predict() — priors propagation
+# ---------------------------------------------------------------------------
+
+class TestPredictPriorsPropagate:
+ """
+ Verify that predict() uses orchestrator.priors (the original calibration
+ priors, which may contain pinned hyperpriors) when calling Predictive,
+ not new_orchestrator.priors (the copy's default/reset priors).
+
+ Regression test for the bug where copy_orchestrator() always creates
+ priors with default values (e.g. theta_values=[[0.5]]) which would be
+ used instead of the fitted calibration priors.
+ """
+
+ def test_predict_passes_original_priors_to_predictive(self, dummy_orchestrator, mocker):
+ """
+ predict() must forward orchestrator.priors (not copy's priors) to
+ the Predictive call. We patch Predictive to capture what priors
+ argument was passed and assert it is the ORIGINAL orchestrator's
+ priors object.
+ """
+ original_priors = dummy_orchestrator.priors
+ # Sanity-check: copy_orchestrator creates an orchestrator whose priors
+ # differ from the original (defaults are reset on copy).
+ copy_orch = copy_orchestrator(dummy_orchestrator)
+ # Priors are not the same object (copy builds fresh defaults).
+ assert copy_orch.priors is not original_priors
+
+ # Capture the priors argument that reach Predictive.__call__.
+ captured = {}
+
+ real_predictive_class = __import__(
+ "numpyro.infer", fromlist=["Predictive"]
+ ).Predictive
+
+ class CapturingPredictive:
+ """Thin wrapper that records the priors kwarg then delegates."""
+ def __init__(self, *args, **kwargs):
+ self._inner = real_predictive_class(*args, **kwargs)
+
+ def __call__(self, rng_key, **kwargs):
+ captured["priors"] = kwargs.get("priors")
+ return self._inner(rng_key, **kwargs)
+
+ mocker.patch(
+ "tfscreen.tfmodel.analysis.prediction.Predictive",
+ CapturingPredictive,
+ )
+
+ # Build a trivial 1-sample posterior dict (MAP-style, no _auto_loc).
+ # Just need the latent sites that exist in the dummy model trace.
+ from numpyro.handlers import seed, trace as nptrace
+ seeded = seed(dummy_orchestrator.jax_model, rng_seed=0)
+ model_tr = nptrace(seeded).get_trace(
+ data=dummy_orchestrator.data,
+ priors=dummy_orchestrator.priors,
+ )
+ fake_posteriors = {
+ name: np.zeros((1,) + site["value"].shape)
+ for name, site in model_tr.items()
+ if site["type"] == "sample" and not site.get("is_observed", False)
+ }
+
+ predict(
+ dummy_orchestrator,
+ fake_posteriors,
+ predict_sites=["growth_pred"],
+ num_samples=None,
+ )
+
+ assert "priors" in captured, "CapturingPredictive was never called"
+ assert captured["priors"] is original_priors, (
+ "predict() passed the copy's priors to Predictive instead of "
+ "the original orchestrator's priors. This regression means "
+ "pinned hyperpriors (e.g. activity_hyper_loc/scale in the "
+ "calibration model) will be randomly re-sampled from the "
+ "default prior, producing wildly incorrect predictions."
+ )
diff --git a/tests/tfscreen/tfmodel/test_presplit.py b/tests/tfscreen/tfmodel/test_presplit.py
new file mode 100644
index 00000000..edc229a2
--- /dev/null
+++ b/tests/tfscreen/tfmodel/test_presplit.py
@@ -0,0 +1,250 @@
+"""
+Tests for the optional pre-split (t = -t_pre) data path.
+
+Coverage:
+ - _read_presplit_df: column checks, silent genotype dropping
+ - _build_presplit_tm: tensor shape/alignment with growth_tm
+ - ModelOrchestrator: presplit=None (no change), presplit passed in
+ - jax_model: presplit observation site fires when data.presplit is set
+"""
+
+import pytest
+import numpy as np
+import pandas as pd
+import jax.numpy as jnp
+import numpyro
+from numpyro.handlers import trace, seed
+
+from tfscreen.tfmodel.model_orchestrator import (
+ _read_presplit_df,
+ _build_presplit_tm,
+ _read_growth_df,
+ _build_growth_tm,
+)
+from tfscreen.tfmodel.data_class import PreSplitData, DataClass
+from tfscreen.tfmodel.tensors.populate_dataclass import populate_dataclass
+
+
+# ---------------------------------------------------------------------------
+# Shared fixtures
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def minimal_growth_df():
+ """Minimal growth DataFrame with two genotypes, two replicates, two
+ condition_pre values — enough to exercise the presplit alignment logic."""
+ rows = []
+ for rep in [1, 2]:
+ for cp in ["kanR-cond", "pheS-cond"]:
+ for geno in ["wt", "A1V", "A2V"]:
+ for t_sel in [60.0, 90.0]:
+ rows.append({
+ "replicate": rep,
+ "condition_pre": cp,
+ "condition_sel": cp,
+ "titrant_name": "iptg",
+ "titrant_conc": 0.0,
+ "t_pre": 30.0,
+ "t_sel": t_sel,
+ "genotype": geno,
+ "ln_cfu": 10.0,
+ "ln_cfu_std": 0.5,
+ })
+ return pd.DataFrame(rows)
+
+
+@pytest.fixture
+def growth_tm(minimal_growth_df):
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ return _build_growth_tm(gdf)
+
+
+@pytest.fixture
+def minimal_presplit_df():
+ """Pre-split data for all (replicate, condition_pre, genotype) combos."""
+ rows = []
+ for rep in [1, 2]:
+ for cp in ["kanR-cond", "pheS-cond"]:
+ for geno in ["wt", "A1V", "A2V"]:
+ rows.append({
+ "replicate": rep,
+ "condition_pre": cp,
+ "genotype": geno,
+ "ln_cfu": 9.5,
+ "ln_cfu_std": 0.4,
+ })
+ return pd.DataFrame(rows)
+
+
+# ---------------------------------------------------------------------------
+# _read_presplit_df
+# ---------------------------------------------------------------------------
+
+def test_read_presplit_df_passes_through_valid(minimal_presplit_df,
+ minimal_growth_df):
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ result = _read_presplit_df(minimal_presplit_df, gdf)
+ assert len(result) == len(minimal_presplit_df)
+ for col in ["replicate", "condition_pre", "genotype", "ln_cfu", "ln_cfu_std"]:
+ assert col in result.columns
+
+
+def test_read_presplit_df_drops_unknown_genotypes(minimal_presplit_df,
+ minimal_growth_df, capsys):
+ df_extra = minimal_presplit_df.copy()
+ extra_row = pd.DataFrame([{"replicate": 1, "condition_pre": "kanR-cond",
+ "genotype": "A99V",
+ "ln_cfu": 5.0, "ln_cfu_std": 0.1}])
+ df_extra = pd.concat([df_extra, extra_row], ignore_index=True)
+
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ result = _read_presplit_df(df_extra, gdf)
+ captured = capsys.readouterr()
+ assert "will be dropped" in captured.out
+ assert "A99V" in captured.out
+
+ assert "A99V" not in result["genotype"].values
+ assert len(result) == len(minimal_presplit_df)
+
+
+def test_read_presplit_df_missing_column_raises(minimal_growth_df):
+ bad_df = pd.DataFrame({
+ "replicate": [1], "condition_pre": ["kanR-cond"],
+ "genotype": ["wt"], "ln_cfu": [9.0]
+ # missing ln_cfu_std
+ })
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ with pytest.raises(Exception):
+ _read_presplit_df(bad_df, gdf)
+
+
+# ---------------------------------------------------------------------------
+# _build_presplit_tm
+# ---------------------------------------------------------------------------
+
+def test_build_presplit_tm_shape(minimal_presplit_df, minimal_growth_df,
+ growth_tm):
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ psdf = _read_presplit_df(minimal_presplit_df, gdf)
+ ps_tm = _build_presplit_tm(psdf, growth_tm)
+
+ # shape should be (num_replicate, num_condition_pre, num_genotype)
+ assert ps_tm.tensor_shape == (2, 2, 3)
+ assert "ln_cfu" in ps_tm.tensors
+ assert "ln_cfu_std" in ps_tm.tensors
+ assert "good_mask" in ps_tm.tensors
+
+
+def test_build_presplit_tm_genotype_alignment(minimal_presplit_df,
+ minimal_growth_df, growth_tm):
+ """Genotype ordering in presplit_tm must match growth_tm."""
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ psdf = _read_presplit_df(minimal_presplit_df, gdf)
+ ps_tm = _build_presplit_tm(psdf, growth_tm)
+
+ growth_geno_labels = growth_tm.tensor_dim_labels[
+ growth_tm.tensor_dim_names.index("genotype")
+ ]
+ presplit_geno_labels = ps_tm.tensor_dim_labels[
+ ps_tm.tensor_dim_names.index("genotype")
+ ]
+ assert list(growth_geno_labels) == list(presplit_geno_labels)
+
+
+def test_build_presplit_tm_partial_coverage(minimal_growth_df, growth_tm):
+ """Genotypes absent from presplit_df get NaN (masked) in the tensor."""
+ partial_df = pd.DataFrame([
+ {"replicate": 1, "condition_pre": "kanR-cond", "genotype": "wt",
+ "ln_cfu": 9.5, "ln_cfu_std": 0.4},
+ ])
+ gdf = _read_growth_df(minimal_growth_df.copy())
+ psdf = _read_presplit_df(partial_df, gdf)
+ ps_tm = _build_presplit_tm(psdf, growth_tm)
+
+ # Only one entry should be valid; the rest masked
+ assert ps_tm.tensors["good_mask"].sum() == 1
+
+
+# ---------------------------------------------------------------------------
+# ModelOrchestrator integration (light smoke test, no JAX tracing)
+# ---------------------------------------------------------------------------
+
+def test_model_orchestrator_without_presplit(minimal_growth_df,
+ minimal_presplit_df):
+ """Passing presplit_df=None leaves data.presplit as None."""
+ from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "A1V", "A2V"],
+ "titrant_name": ["iptg"] * 3,
+ "titrant_conc": [0.0] * 3,
+ "theta_obs": [0.5] * 3,
+ "theta_std": [0.1] * 3,
+ })
+ orchestrator = ModelOrchestrator(minimal_growth_df, binding_df, presplit_df=None)
+ assert orchestrator.data.presplit is None
+
+
+def test_model_orchestrator_with_presplit(minimal_growth_df,
+ minimal_presplit_df):
+ """Passing presplit_df populates data.presplit as a PreSplitData."""
+ from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "A1V", "A2V"],
+ "titrant_name": ["iptg"] * 3,
+ "titrant_conc": [0.0] * 3,
+ "theta_obs": [0.5] * 3,
+ "theta_std": [0.1] * 3,
+ })
+ orchestrator = ModelOrchestrator(minimal_growth_df, binding_df,
+ presplit_df=minimal_presplit_df)
+ ps = orchestrator.data.presplit
+ assert ps is not None
+ assert isinstance(ps, PreSplitData)
+ assert ps.num_replicate == 2
+ assert ps.num_condition_pre == 2
+ assert ps.num_genotype == 3
+ assert ps.ln_cfu_t0.shape == (2, 2, 3)
+ assert ps.good_mask.shape == (2, 2, 3)
+ # All entries should be valid (full coverage)
+ assert bool(jnp.all(ps.good_mask))
+
+
+# ---------------------------------------------------------------------------
+# jax_model: presplit_obs site fires
+# ---------------------------------------------------------------------------
+
+def test_jax_model_presplit_obs_site_present(minimal_growth_df,
+ minimal_presplit_df):
+ """The presplit_obs sample site should appear in the model trace when
+ data.presplit is not None."""
+ from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "A1V", "A2V"],
+ "titrant_name": ["iptg"] * 3,
+ "titrant_conc": [0.0] * 3,
+ "theta_obs": [0.5] * 3,
+ "theta_std": [0.1] * 3,
+ })
+ orchestrator = ModelOrchestrator(minimal_growth_df, binding_df,
+ presplit_df=minimal_presplit_df)
+
+ model_fn = orchestrator.jax_model
+ tr = trace(seed(model_fn, 0)).get_trace(orchestrator.data, orchestrator.priors)
+ assert "presplit_obs" in tr
+
+
+def test_jax_model_no_presplit_obs_site_without_data(minimal_growth_df):
+ """The presplit_obs site must NOT appear when presplit_df is not given."""
+ from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
+ binding_df = pd.DataFrame({
+ "genotype": ["wt", "A1V", "A2V"],
+ "titrant_name": ["iptg"] * 3,
+ "titrant_conc": [0.0] * 3,
+ "theta_obs": [0.5] * 3,
+ "theta_std": [0.1] * 3,
+ })
+ orchestrator = ModelOrchestrator(minimal_growth_df, binding_df)
+
+ model_fn = orchestrator.jax_model
+ tr = trace(seed(model_fn, 0)).get_trace(orchestrator.data, orchestrator.priors)
+ assert "presplit_obs" not in tr
diff --git a/tests/tfscreen/tfmodel/test_prior_predictive.py b/tests/tfscreen/tfmodel/test_prior_predictive.py
index 53913d49..5ba65a11 100644
--- a/tests/tfscreen/tfmodel/test_prior_predictive.py
+++ b/tests/tfscreen/tfmodel/test_prior_predictive.py
@@ -7,13 +7,13 @@
@pytest.fixture
-def dummy_mc():
+def dummy_orchestrator():
growth_df = pd.DataFrame({
"genotype": ["wt", "wt", "M42V", "M42V"],
"titrant_name": ["tit1", "tit1", "tit1", "tit1"],
"titrant_conc": [0.0, 1.0, 0.0, 1.0],
- "condition_pre": ["pre1", "pre1", "pre1", "pre1"],
- "condition_sel": ["sel1", "sel1", "sel1", "sel1"],
+ "condition_pre": ["pre-1", "pre-1", "pre-1", "pre-1"],
+ "condition_sel": ["sel+1", "sel+1", "sel+1", "sel+1"],
"t_pre": [10.0, 10.0, 10.0, 10.0],
"t_sel": [0.0, 20.0, 0.0, 20.0],
"ln_cfu": [0.0, 5.0, 0.0, 3.0],
@@ -34,54 +34,54 @@ def dummy_mc():
# draw_prior
# ---------------------------------------------------------------------------
-def test_draw_prior_returns_dicts(dummy_mc):
- predictions, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
+def test_draw_prior_returns_dicts(dummy_orchestrator):
+ predictions, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
assert isinstance(predictions, dict)
assert isinstance(latent_params, dict)
-def test_draw_prior_contains_growth_pred(dummy_mc):
- predictions, _ = draw_prior(dummy_mc, rng_key=0, num_draws=1)
+def test_draw_prior_contains_growth_pred(dummy_orchestrator):
+ predictions, _ = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
assert "growth_pred" in predictions
-def test_draw_prior_growth_pred_shape(dummy_mc):
- predictions, _ = draw_prior(dummy_mc, rng_key=0, num_draws=3)
+def test_draw_prior_growth_pred_shape(dummy_orchestrator):
+ predictions, _ = draw_prior(dummy_orchestrator, rng_key=0, num_draws=3)
gp = predictions["growth_pred"]
assert gp.shape[0] == 3, "First dim should equal num_draws"
assert gp.ndim > 1, "growth_pred should have spatial dims beyond num_draws"
-def test_draw_prior_latent_params_shape(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=3)
+def test_draw_prior_latent_params_shape(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=3)
for name, val in latent_params.items():
assert val.shape[0] == 3, f"Parameter '{name}' first dim should equal num_draws"
-def test_draw_prior_no_observed_in_latent(dummy_mc):
+def test_draw_prior_no_observed_in_latent(dummy_orchestrator):
# Observed sites (obs= in model) should not appear in latent_params.
# The growth observation site is named 'final_binding_obs_growth_obs'.
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
assert "final_binding_obs_growth_obs" not in latent_params
-def test_draw_prior_reproducible(dummy_mc):
- _, p1 = draw_prior(dummy_mc, rng_key=42, num_draws=1)
- _, p2 = draw_prior(dummy_mc, rng_key=42, num_draws=1)
+def test_draw_prior_reproducible(dummy_orchestrator):
+ _, p1 = draw_prior(dummy_orchestrator, rng_key=42, num_draws=1)
+ _, p2 = draw_prior(dummy_orchestrator, rng_key=42, num_draws=1)
for k in p1:
np.testing.assert_array_equal(p1[k], p2[k])
-def test_draw_prior_different_seeds(dummy_mc):
- _, p1 = draw_prior(dummy_mc, rng_key=0, num_draws=1)
- _, p2 = draw_prior(dummy_mc, rng_key=1, num_draws=1)
+def test_draw_prior_different_seeds(dummy_orchestrator):
+ _, p1 = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
+ _, p2 = draw_prior(dummy_orchestrator, rng_key=1, num_draws=1)
# At least one parameter should differ across seeds.
any_diff = any(not np.array_equal(p1[k], p2[k]) for k in p1 if k in p2)
assert any_diff
-def test_draw_prior_multiple_independent_draws(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=5)
+def test_draw_prior_multiple_independent_draws(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=5)
# Draws should not all be identical.
first_key = next(iter(latent_params))
v = latent_params[first_key]
@@ -93,54 +93,54 @@ def test_draw_prior_multiple_independent_draws(dummy_mc):
# growth_df_from_prior
# ---------------------------------------------------------------------------
-def test_growth_df_from_prior_returns_dataframe(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
- df = growth_df_from_prior(dummy_mc, latent_params)
+def test_growth_df_from_prior_returns_dataframe(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
+ df = growth_df_from_prior(dummy_orchestrator, latent_params)
assert isinstance(df, pd.DataFrame)
-def test_growth_df_from_prior_same_rows(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
- df = growth_df_from_prior(dummy_mc, latent_params)
- assert len(df) == len(dummy_mc.growth_df)
+def test_growth_df_from_prior_same_rows(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
+ df = growth_df_from_prior(dummy_orchestrator, latent_params)
+ assert len(df) == len(dummy_orchestrator.growth_df)
-def test_growth_df_from_prior_has_ln_cfu(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
- df = growth_df_from_prior(dummy_mc, latent_params)
+def test_growth_df_from_prior_has_ln_cfu(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
+ df = growth_df_from_prior(dummy_orchestrator, latent_params)
assert "ln_cfu" in df.columns
assert df["ln_cfu"].notna().all()
-def test_growth_df_from_prior_preserves_metadata(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
- df = growth_df_from_prior(dummy_mc, latent_params)
+def test_growth_df_from_prior_preserves_metadata(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
+ df = growth_df_from_prior(dummy_orchestrator, latent_params)
for col in ["genotype", "titrant_conc", "ln_cfu_std", "replicate"]:
assert col in df.columns
-def test_growth_df_from_prior_noise_changes_ln_cfu(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=1)
- df_clean = growth_df_from_prior(dummy_mc, latent_params, noise_rng=None)
+def test_growth_df_from_prior_noise_changes_ln_cfu(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=1)
+ df_clean = growth_df_from_prior(dummy_orchestrator, latent_params, noise_rng=None)
rng = np.random.default_rng(7)
- df_noisy = growth_df_from_prior(dummy_mc, latent_params, noise_rng=rng)
+ df_noisy = growth_df_from_prior(dummy_orchestrator, latent_params, noise_rng=rng)
# Noisy and clean should differ (adding N(0, 0.1) noise almost certainly changes values).
assert not np.allclose(df_clean["ln_cfu"].values, df_noisy["ln_cfu"].values)
-def test_growth_df_from_prior_different_draws(dummy_mc):
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=5)
- df0 = growth_df_from_prior(dummy_mc, latent_params, draw_idx=0)
- df1 = growth_df_from_prior(dummy_mc, latent_params, draw_idx=1)
+def test_growth_df_from_prior_different_draws(dummy_orchestrator):
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=5)
+ df0 = growth_df_from_prior(dummy_orchestrator, latent_params, draw_idx=0)
+ df1 = growth_df_from_prior(dummy_orchestrator, latent_params, draw_idx=1)
# Different draws should generally produce different predictions.
assert not np.allclose(df0["ln_cfu"].values, df1["ln_cfu"].values)
-def test_latent_params_compatible_with_predict(dummy_mc):
+def test_latent_params_compatible_with_predict(dummy_orchestrator):
"""latent_params from draw_prior can be passed to predict() without error."""
from tfscreen.tfmodel.analysis.prediction import predict
- _, latent_params = draw_prior(dummy_mc, rng_key=0, num_draws=2)
- result = predict(dummy_mc, latent_params, predict_sites=["growth_pred"],
+ _, latent_params = draw_prior(dummy_orchestrator, rng_key=0, num_draws=2)
+ result = predict(dummy_orchestrator, latent_params, predict_sites=["growth_pred"],
num_samples=None, num_marginal_samples=2)
- assert "median" in result.columns or any(c in result.columns
- for c in ["lower_95", "upper_95"])
+ assert "q0.5" in result.columns or any(c in result.columns
+ for c in ["q0.025", "q0.975"])
diff --git a/tests/tfscreen/tfmodel/test_sbc.py b/tests/tfscreen/tfmodel/test_sbc.py
deleted file mode 100644
index a69bd718..00000000
--- a/tests/tfscreen/tfmodel/test_sbc.py
+++ /dev/null
@@ -1,284 +0,0 @@
-"""
-Tests for tfscreen.tfmodel.analysis.sbc.
-"""
-
-import os
-import tempfile
-
-import h5py
-import numpy as np
-import pytest
-
-from tfscreen.tfmodel.analysis.sbc import (
- _find_pairs,
- _load_h5_params,
- compute_sbc_ranks,
- summarize_sbc,
-)
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-def _write_h5(path, arrays):
- """Write a dict of numpy arrays to an HDF5 file."""
- with h5py.File(path, "w") as hf:
- for key, val in arrays.items():
- hf.create_dataset(key, data=val)
-
-
-# ---------------------------------------------------------------------------
-# _load_h5_params
-# ---------------------------------------------------------------------------
-
-class TestLoadH5Params:
-
- def test_round_trips_arrays(self, tmp_path):
- data = {"alpha": np.array([1.0, 2.0, 3.0]), "beta": np.array([[4.0, 5.0]])}
- p = str(tmp_path / "test.h5")
- _write_h5(p, data)
- result = _load_h5_params(p)
- np.testing.assert_allclose(result["alpha"], data["alpha"])
- np.testing.assert_allclose(result["beta"], data["beta"])
-
- def test_returns_all_keys(self, tmp_path):
- data = {"x": np.ones(3), "y": np.zeros(5)}
- p = str(tmp_path / "test.h5")
- _write_h5(p, data)
- result = _load_h5_params(p)
- assert set(result.keys()) == {"x", "y"}
-
- def test_empty_file(self, tmp_path):
- p = str(tmp_path / "empty.h5")
- _write_h5(p, {})
- result = _load_h5_params(p)
- assert result == {}
-
-
-# ---------------------------------------------------------------------------
-# compute_sbc_ranks
-# ---------------------------------------------------------------------------
-
-class TestComputeSBCRanks:
-
- def _make_pair(self, tmp_path, gt_val, posterior_samples, name="param"):
- """Write ground-truth (shape (1, *)) and posterior (shape (S, *)) files."""
- gt_path = str(tmp_path / "run_ground_truth.h5")
- post_path = str(tmp_path / "run_posterior.h5")
- _write_h5(gt_path, {name: np.array(gt_val)[np.newaxis, :]})
- _write_h5(post_path, {name: np.array(posterior_samples)})
- return gt_path, post_path
-
- def test_rank_zero_when_gt_below_all_posterior(self, tmp_path):
- # gt = -100; all posterior > -100 → rank = 0
- gt_path, post_path = self._make_pair(
- tmp_path,
- gt_val=[-100.0],
- posterior_samples=np.ones((20, 1)),
- )
- ranks = compute_sbc_ranks(gt_path, post_path)
- assert "param" in ranks
- np.testing.assert_allclose(ranks["param"], [0.0])
-
- def test_rank_one_when_gt_above_all_posterior(self, tmp_path):
- # gt = 100; all posterior < 100 → rank = 1
- gt_path, post_path = self._make_pair(
- tmp_path,
- gt_val=[100.0],
- posterior_samples=-np.ones((20, 1)),
- )
- ranks = compute_sbc_ranks(gt_path, post_path)
- np.testing.assert_allclose(ranks["param"], [1.0])
-
- def test_rank_half_when_gt_at_median(self, tmp_path):
- # posterior is [-10, -9, ..., -1, 0, 1, ..., 10] (21 values), gt = 0
- # 10 values < 0, 11 >= 0 → rank = 10/21
- post = np.arange(-10, 11, dtype=float).reshape(21, 1)
- gt_path, post_path = self._make_pair(
- tmp_path, gt_val=[0.0], posterior_samples=post
- )
- ranks = compute_sbc_ranks(gt_path, post_path)
- np.testing.assert_allclose(ranks["param"], [10 / 21])
-
- def test_multidimensional_param(self, tmp_path):
- # param shape (3,): three independent elements
- rng = np.random.default_rng(0)
- post = rng.normal(size=(100, 3))
- gt = np.zeros((1, 3))
- gt_path = str(tmp_path / "run_ground_truth.h5")
- post_path = str(tmp_path / "run_posterior.h5")
- _write_h5(gt_path, {"param": gt})
- _write_h5(post_path, {"param": post})
- ranks = compute_sbc_ranks(gt_path, post_path)
- assert ranks["param"].shape == (3,)
- assert np.all((ranks["param"] >= 0) & (ranks["param"] <= 1))
-
- def test_missing_posterior_param_skipped(self, tmp_path):
- gt_path = str(tmp_path / "run_ground_truth.h5")
- post_path = str(tmp_path / "run_posterior.h5")
- _write_h5(gt_path, {"param_a": np.ones((1, 2)), "param_b": np.ones((1, 2))})
- _write_h5(post_path, {"param_a": np.ones((10, 2))})
- ranks = compute_sbc_ranks(gt_path, post_path)
- assert "param_a" in ranks
- assert "param_b" not in ranks
-
- def test_returns_empty_when_no_common_params(self, tmp_path):
- gt_path = str(tmp_path / "run_ground_truth.h5")
- post_path = str(tmp_path / "run_posterior.h5")
- _write_h5(gt_path, {"a": np.ones((1, 2))})
- _write_h5(post_path, {"b": np.ones((10, 2))})
- ranks = compute_sbc_ranks(gt_path, post_path)
- assert ranks == {}
-
- def test_rank_values_in_unit_interval(self, tmp_path):
- rng = np.random.default_rng(42)
- post = rng.normal(size=(50, 4))
- gt = rng.normal(size=(1, 4))
- gt_path = str(tmp_path / "run_ground_truth.h5")
- post_path = str(tmp_path / "run_posterior.h5")
- _write_h5(gt_path, {"theta": gt})
- _write_h5(post_path, {"theta": post})
- ranks = compute_sbc_ranks(gt_path, post_path)
- assert np.all((ranks["theta"] >= 0) & (ranks["theta"] <= 1))
-
-
-# ---------------------------------------------------------------------------
-# _find_pairs
-# ---------------------------------------------------------------------------
-
-class TestFindPairs:
-
- def test_finds_paired_files(self, tmp_path):
- for prefix in ("run001", "run002"):
- (tmp_path / f"{prefix}_ground_truth.h5").touch()
- (tmp_path / f"{prefix}_posterior.h5").touch()
- pairs = _find_pairs(str(tmp_path))
- assert len(pairs) == 2
- run_ids = {p[0] for p in pairs}
- assert run_ids == {"run001", "run002"}
-
- def test_missing_posterior_gives_none(self, tmp_path):
- (tmp_path / "run001_ground_truth.h5").touch()
- pairs = _find_pairs(str(tmp_path))
- assert len(pairs) == 1
- run_id, gt_path, post_path = pairs[0]
- assert run_id == "run001"
- assert post_path is None
-
- def test_returns_empty_when_no_gt_files(self, tmp_path):
- (tmp_path / "run001_posterior.h5").touch()
- pairs = _find_pairs(str(tmp_path))
- assert pairs == []
-
- def test_gt_path_is_absolute(self, tmp_path):
- (tmp_path / "run001_ground_truth.h5").touch()
- (tmp_path / "run001_posterior.h5").touch()
- pairs = _find_pairs(str(tmp_path))
- _, gt_path, post_path = pairs[0]
- assert os.path.isabs(gt_path)
- assert os.path.isabs(post_path)
-
-
-# ---------------------------------------------------------------------------
-# summarize_sbc (integration)
-# ---------------------------------------------------------------------------
-
-def _make_calibrated_sbc_dir(tmp_path, n_runs=20, n_samples=200, seed=0):
- """
- Create a synthetic SBC directory where ranks are genuinely uniform.
-
- For each run: draw gt ~ N(0,1), draw posterior ~ N(0,1) (S samples).
- The ranks will be approximately uniform under this prior-predictive setup.
- """
- rng = np.random.default_rng(seed)
- sbc_dir = tmp_path / "sbc"
- sbc_dir.mkdir()
- for i in range(n_runs):
- gt = rng.normal(size=(1, 3))
- post = rng.normal(size=(n_samples, 3))
- run_id = f"run{i:04d}"
- _write_h5(str(sbc_dir / f"{run_id}_ground_truth.h5"), {"alpha": gt})
- _write_h5(str(sbc_dir / f"{run_id}_posterior.h5"), {"alpha": post})
- return str(sbc_dir)
-
-
-class TestSummarizeSBC:
-
- def test_returns_dataframe_with_expected_columns(self, tmp_path):
- sbc_dir = _make_calibrated_sbc_dir(tmp_path)
- df = summarize_sbc(sbc_dir)
- assert not df.empty
- for col in ("param", "n_runs", "n_ranks", "mean_rank", "ks_stat", "ks_pval"):
- assert col in df.columns
-
- def test_one_row_per_parameter(self, tmp_path):
- sbc_dir = _make_calibrated_sbc_dir(tmp_path)
- df = summarize_sbc(sbc_dir)
- assert len(df) == 1 # only "alpha"
- assert df.iloc[0]["param"] == "alpha"
-
- def test_n_runs_correct(self, tmp_path):
- n_runs = 15
- sbc_dir = _make_calibrated_sbc_dir(tmp_path, n_runs=n_runs)
- df = summarize_sbc(sbc_dir)
- assert df.iloc[0]["n_runs"] == n_runs
-
- def test_n_ranks_correct(self, tmp_path):
- # 10 runs × 3 elements each = 30 rank values
- sbc_dir = _make_calibrated_sbc_dir(tmp_path, n_runs=10)
- df = summarize_sbc(sbc_dir)
- assert df.iloc[0]["n_ranks"] == 30
-
- def test_writes_summary_csv(self, tmp_path):
- sbc_dir = _make_calibrated_sbc_dir(tmp_path)
- summarize_sbc(sbc_dir)
- assert os.path.exists(os.path.join(sbc_dir, "sbc_sbc_summary.csv"))
-
- def test_writes_ranks_csv(self, tmp_path):
- sbc_dir = _make_calibrated_sbc_dir(tmp_path)
- summarize_sbc(sbc_dir)
- assert os.path.exists(os.path.join(sbc_dir, "sbc_sbc_ranks.csv"))
-
- def test_writes_histogram_pdf(self, tmp_path):
- sbc_dir = _make_calibrated_sbc_dir(tmp_path)
- summarize_sbc(sbc_dir)
- assert os.path.exists(os.path.join(sbc_dir, "sbc_rank_hist.pdf"))
-
- def test_custom_out_prefix(self, tmp_path):
- sbc_dir = _make_calibrated_sbc_dir(tmp_path)
- prefix = str(tmp_path / "myprefix")
- summarize_sbc(sbc_dir, out_prefix=prefix)
- assert os.path.exists(f"{prefix}_sbc_summary.csv")
- assert os.path.exists(f"{prefix}_sbc_ranks.csv")
-
- def test_empty_dir_returns_empty_dataframe(self, tmp_path):
- sbc_dir = str(tmp_path / "empty")
- os.makedirs(sbc_dir)
- df = summarize_sbc(sbc_dir)
- assert df.empty
-
- def test_missing_posterior_skipped_without_crash(self, tmp_path):
- sbc_dir = tmp_path / "sbc"
- sbc_dir.mkdir()
- rng = np.random.default_rng(1)
- # Run 0: has both files
- gt = rng.normal(size=(1, 2))
- post = rng.normal(size=(50, 2))
- _write_h5(str(sbc_dir / "run0000_ground_truth.h5"), {"alpha": gt})
- _write_h5(str(sbc_dir / "run0000_posterior.h5"), {"alpha": post})
- # Run 1: only ground truth
- _write_h5(str(sbc_dir / "run0001_ground_truth.h5"), {"alpha": rng.normal(size=(1, 2))})
- df = summarize_sbc(str(sbc_dir))
- assert not df.empty
- assert df.iloc[0]["n_runs"] == 1
-
- def test_nonexistent_dir_raises(self, tmp_path):
- with pytest.raises(FileNotFoundError):
- summarize_sbc(str(tmp_path / "does_not_exist"))
-
- def test_mean_rank_near_half_for_calibrated_model(self, tmp_path):
- # With many runs, mean rank should be close to 0.5
- sbc_dir = _make_calibrated_sbc_dir(tmp_path, n_runs=100, n_samples=500)
- df = summarize_sbc(sbc_dir)
- assert abs(df.iloc[0]["mean_rank"] - 0.5) < 0.05
diff --git a/tests/tfscreen/tfmodel/test_share_replicates.py b/tests/tfscreen/tfmodel/test_share_replicates.py
index c10fac5a..20b59c6b 100644
--- a/tests/tfscreen/tfmodel/test_share_replicates.py
+++ b/tests/tfscreen/tfmodel/test_share_replicates.py
@@ -10,8 +10,8 @@ def get_mock_df():
# Let's say we have genotypes A, B
for rep in [1, 2]:
- for cond_pre in ['pre1', 'pre2']:
- for cond_sel in ['sel1']:
+ for cond_pre in ['pre-1', 'pre-2']:
+ for cond_sel in ['sel+1']:
for tname in ['inducer']:
for tconc in [0.0, 1.0]:
for geno in ['wt', 'A1T']:
@@ -32,7 +32,7 @@ def get_mock_df():
df = pd.DataFrame(data)
growth_df = df.copy()
- binding_df = df[df['cond_pre'] == 'pre1'].copy() if False else df.copy() # Just pass identical df for mock
+ binding_df = df[df['cond_pre'] == 'pre-1'].copy() if False else df.copy() # Just pass identical df for mock
return growth_df, binding_df
@@ -60,7 +60,7 @@ def test_growth_shares_replicates():
df_mapped = model_shared.growth_tm.df
# Find rows for rep 1 and rep 2 for the same condition
- rep1_map = df_mapped[(df_mapped["replicate"] == 1) & (df_mapped["condition_pre"] == 'pre1')]["map_condition_pre"].iloc[0]
- rep2_map = df_mapped[(df_mapped["replicate"] == 2) & (df_mapped["condition_pre"] == 'pre1')]["map_condition_pre"].iloc[0]
+ rep1_map = df_mapped[(df_mapped["replicate"] == 1) & (df_mapped["condition_pre"] == 'pre-1')]["map_condition_pre"].iloc[0]
+ rep2_map = df_mapped[(df_mapped["replicate"] == 2) & (df_mapped["condition_pre"] == 'pre-1')]["map_condition_pre"].iloc[0]
assert rep1_map == rep2_map, "Replicates did not receive the same parameter mapping index!"
diff --git a/tests/tfscreen/tfmodel/test_spiked_genotypes.py b/tests/tfscreen/tfmodel/test_spiked_genotypes.py
index 35836413..d36591f9 100644
--- a/tests/tfscreen/tfmodel/test_spiked_genotypes.py
+++ b/tests/tfscreen/tfmodel/test_spiked_genotypes.py
@@ -18,8 +18,8 @@ def dummy_data():
"genotype": g,
"ln_cfu": 1.0,
"ln_cfu_std": 0.1,
- "condition_pre": "pre",
- "condition_sel": "sel",
+ "condition_pre": "pre-cond",
+ "condition_sel": "sel+cond",
"t_pre": 1.0,
"t_sel": 10.0,
"titrant_name": "iptg",
@@ -48,14 +48,14 @@ def test_spiked_genotypes_masking(dummy_data):
# but ModelOrchestrator.__init__ calls _initialize_data which does most of the heavy lifting.
# We can use real data for this unit test since it's small.
- gm = ModelOrchestrator(growth_df, binding_df, spiked_genotypes=spiked)
+ orchestrator = ModelOrchestrator(growth_df, binding_df, spiked_genotypes=spiked)
# Get genotype labels from the growth tensor manager
- genotype_idx = gm.growth_tm.tensor_dim_names.index("genotype")
- genotype_labels = gm.growth_tm.tensor_dim_labels[genotype_idx]
+ genotype_idx = orchestrator.growth_tm.tensor_dim_names.index("genotype")
+ genotype_labels = orchestrator.growth_tm.tensor_dim_labels[genotype_idx]
# Check the congression_mask in the data object
- mask = np.array(gm.data.growth.congression_mask)
+ mask = np.array(orchestrator.data.growth.congression_mask)
for i, label in enumerate(genotype_labels):
if label == "A10G":
assert not mask[i], f"Genotype {label} should be masked (False)"
@@ -87,12 +87,12 @@ def test_setup_batching_zero_division():
def test_spiked_genotypes_single_string(dummy_data):
"""Test that a single string is handled correctly as spiked_genotypes."""
growth_df, binding_df = dummy_data
- gm = ModelOrchestrator(growth_df, binding_df, spiked_genotypes="A10G")
+ orchestrator = ModelOrchestrator(growth_df, binding_df, spiked_genotypes="A10G")
- genotype_idx = gm.growth_tm.tensor_dim_names.index("genotype")
- genotype_labels = gm.growth_tm.tensor_dim_labels[genotype_idx]
+ genotype_idx = orchestrator.growth_tm.tensor_dim_names.index("genotype")
+ genotype_labels = orchestrator.growth_tm.tensor_dim_labels[genotype_idx]
- mask = np.array(gm.data.growth.congression_mask)
+ mask = np.array(orchestrator.data.growth.congression_mask)
for i, label in enumerate(genotype_labels):
if label == "A10G":
assert not mask[i]
diff --git a/tests/tfscreen/tfmodel/test_struct_wiring.py b/tests/tfscreen/tfmodel/test_struct_wiring.py
index 6d25374d..a342bcbd 100644
--- a/tests/tfscreen/tfmodel/test_struct_wiring.py
+++ b/tests/tfscreen/tfmodel/test_struct_wiring.py
@@ -80,9 +80,9 @@ def _make_growth_csv(tmp_path, genotypes=_GENOTYPES):
"num_muts": num_muts,
"counts": 1000,
"replicate": rep,
- "condition_pre": "pre",
+ "condition_pre": "pre-cond",
"t_pre": 24.0,
- "condition_sel": "sel",
+ "condition_sel": "sel+cond",
"t_sel": 48.0,
"titrant_name": "iptg",
"titrant_conc": conc,
@@ -148,13 +148,13 @@ def test_settings_includes_thermo_data(tmp_path):
h5_path = _make_struct_ensemble_h5(tmp_path)
from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
- mc = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
growth_path, binding_path,
theta="thermo.O2_C4_K3_U0_a.PnnC",
thermo_data=h5_path,
)
- assert "thermo_data" in mc.settings
- assert mc.settings["thermo_data"] == h5_path
+ assert "thermo_data" in orchestrator.settings
+ assert orchestrator.settings["thermo_data"] == h5_path
def test_default_thermo_data_is_none(tmp_path):
@@ -163,8 +163,8 @@ def test_default_thermo_data_is_none(tmp_path):
binding_path = _make_binding_csv(tmp_path)
from tfscreen.tfmodel.model_orchestrator import ModelOrchestrator
- mc = ModelOrchestrator(growth_path, binding_path, theta="hill_geno")
- assert mc.settings["thermo_data"] is None
+ orchestrator = ModelOrchestrator(growth_path, binding_path, theta="hill_geno")
+ assert orchestrator.settings["thermo_data"] is None
# ──────────────────────────────────────────────────────────────────────────────
@@ -186,7 +186,7 @@ def test_missing_thermo_data_raises(tmp_path):
# ──────────────────────────────────────────────────────────────────────────────
@pytest.fixture(scope="module")
-def fitted_mc(tmp_path_factory):
+def fitted_orchestrator(tmp_path_factory):
"""
Build a ModelOrchestrator with lac_dimer_lnK_nn_prior and a synthetic HDF5 file.
Expensive to construct, so scoped to the module.
@@ -205,39 +205,39 @@ def fitted_mc(tmp_path_factory):
class TestStructFieldsOnGrowthData:
- def test_struct_names_is_correct_tuple(self, fitted_mc):
- assert fitted_mc.data.growth.struct_names == STRUCTURE_NAMES
+ def test_struct_names_is_correct_tuple(self, fitted_orchestrator):
+ assert fitted_orchestrator.data.growth.struct_names == STRUCTURE_NAMES
- def test_num_struct_is_four(self, fitted_mc):
- assert fitted_mc.data.growth.num_struct == 4
+ def test_num_struct_is_four(self, fitted_orchestrator):
+ assert fitted_orchestrator.data.growth.num_struct == 4
- def test_struct_features_shape(self, fitted_mc):
- M = fitted_mc.data.growth.num_mutation
- S = fitted_mc.data.growth.num_struct
- feat = fitted_mc.data.growth.struct_features
+ def test_struct_features_shape(self, fitted_orchestrator):
+ M = fitted_orchestrator.data.growth.num_mutation
+ S = fitted_orchestrator.data.growth.num_struct
+ feat = fitted_orchestrator.data.growth.struct_features
assert feat.shape == (M, S, 60)
- def test_struct_n_chains_shape(self, fitted_mc):
- S = fitted_mc.data.growth.num_struct
- assert fitted_mc.data.growth.struct_n_chains.shape == (S,)
+ def test_struct_n_chains_shape(self, fitted_orchestrator):
+ S = fitted_orchestrator.data.growth.num_struct
+ assert fitted_orchestrator.data.growth.struct_n_chains.shape == (S,)
- def test_struct_n_chains_value(self, fitted_mc):
+ def test_struct_n_chains_value(self, fitted_orchestrator):
# All structures were written with n_chains_bearing_mut=2
import jax.numpy as jnp
- n = fitted_mc.data.growth.struct_n_chains
+ n = fitted_orchestrator.data.growth.struct_n_chains
assert int(n[0]) == 2
- def test_no_contact_distances_without_epistasis(self, fitted_mc):
+ def test_no_contact_distances_without_epistasis(self, fitted_orchestrator):
# Fixture was built without epistasis=True, so no contact arrays
- assert fitted_mc.data.growth.struct_contact_distances is None
- assert fitted_mc.data.growth.struct_contact_pair_idx is None
+ assert fitted_orchestrator.data.growth.struct_contact_distances is None
+ assert fitted_orchestrator.data.growth.struct_contact_pair_idx is None
- def test_mut_labels_populated(self, fitted_mc):
+ def test_mut_labels_populated(self, fitted_orchestrator):
# _needs_mut must have fired for lac_dimer_lnK_nn_prior
- assert len(fitted_mc.mut_labels) > 0
+ assert len(fitted_orchestrator.mut_labels) > 0
- def test_num_mutation_nonzero(self, fitted_mc):
- assert fitted_mc.data.growth.num_mutation > 0
+ def test_num_mutation_nonzero(self, fitted_orchestrator):
+ assert fitted_orchestrator.data.growth.num_mutation > 0
class TestStructFieldsWithEpistasis:
@@ -288,14 +288,14 @@ def test_svi_runs_without_error(tmp_path):
binding_path = _make_binding_csv(tmp_path)
h5_path = _make_struct_ensemble_h5(tmp_path)
- mc = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
growth_path, binding_path,
theta="thermo.O2_C4_K3_U0_a.PnnC",
thermo_data=h5_path,
)
out_prefix = str(tmp_path / "svi_out")
- inference = RunInference(model=mc, seed=0)
+ inference = RunInference(model=orchestrator, seed=0)
svi = inference.setup_svi(adam_step_size=1e-3)
svi_state, params, converged = inference.run_optimization(
svi=svi,
@@ -319,7 +319,7 @@ def test_svi_with_epistasis_runs_without_error(tmp_path):
binding_path = _make_binding_csv(tmp_path)
h5_path = _make_struct_ensemble_h5(tmp_path)
- mc = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
growth_path, binding_path,
theta="thermo.O2_C4_K3_U0_a.PnnC",
thermo_data=h5_path,
@@ -327,7 +327,7 @@ def test_svi_with_epistasis_runs_without_error(tmp_path):
)
out_prefix = str(tmp_path / "svi_epi_out")
- inference = RunInference(model=mc, seed=0)
+ inference = RunInference(model=orchestrator, seed=0)
svi = inference.setup_svi(adam_step_size=1e-3)
svi_state, params, converged = inference.run_optimization(
svi=svi,
@@ -391,7 +391,7 @@ def test_binding_only_struct_raises_without_thermo_data(tmp_path):
@pytest.fixture(scope="module")
-def binding_only_ddG_mc(tmp_path_factory):
+def binding_only_ddG_orchestrator(tmp_path_factory):
"""
ModelOrchestrator in binding_only mode with mwc_dimer_lnK_ddG_prior.
Scoped to module so it is built once for all tests below.
@@ -412,30 +412,30 @@ def binding_only_ddG_mc(tmp_path_factory):
class TestBindingOnlyDdGStructFields:
"""BindingData must carry the struct_* fields when using mwc_dimer_lnK_ddG_prior."""
- def test_struct_names_is_correct_tuple(self, binding_only_ddG_mc):
- assert set(binding_only_ddG_mc.data.binding.struct_names) == set(_MWC_STRUCTURE_NAMES)
+ def test_struct_names_is_correct_tuple(self, binding_only_ddG_orchestrator):
+ assert set(binding_only_ddG_orchestrator.data.binding.struct_names) == set(_MWC_STRUCTURE_NAMES)
- def test_num_struct_is_six(self, binding_only_ddG_mc):
- assert binding_only_ddG_mc.data.binding.num_struct == 6
+ def test_num_struct_is_six(self, binding_only_ddG_orchestrator):
+ assert binding_only_ddG_orchestrator.data.binding.num_struct == 6
- def test_struct_features_shape(self, binding_only_ddG_mc):
- M = binding_only_ddG_mc.data.binding.num_mutation
- S = binding_only_ddG_mc.data.binding.num_struct
- feat = binding_only_ddG_mc.data.binding.struct_features
+ def test_struct_features_shape(self, binding_only_ddG_orchestrator):
+ M = binding_only_ddG_orchestrator.data.binding.num_mutation
+ S = binding_only_ddG_orchestrator.data.binding.num_struct
+ feat = binding_only_ddG_orchestrator.data.binding.struct_features
assert feat.shape == (M, S)
- def test_mut_labels_populated(self, binding_only_ddG_mc):
- assert len(binding_only_ddG_mc.mut_labels) > 0
+ def test_mut_labels_populated(self, binding_only_ddG_orchestrator):
+ assert len(binding_only_ddG_orchestrator.mut_labels) > 0
- def test_num_mutation_nonzero(self, binding_only_ddG_mc):
- assert binding_only_ddG_mc.data.binding.num_mutation > 0
+ def test_num_mutation_nonzero(self, binding_only_ddG_orchestrator):
+ assert binding_only_ddG_orchestrator.data.binding.num_mutation > 0
- def test_growth_data_is_none(self, binding_only_ddG_mc):
- assert binding_only_ddG_mc.data.growth is None
+ def test_growth_data_is_none(self, binding_only_ddG_orchestrator):
+ assert binding_only_ddG_orchestrator.data.growth is None
- def test_struct_n_chains_is_none_for_ddG_csv(self, binding_only_ddG_mc):
+ def test_struct_n_chains_is_none_for_ddG_csv(self, binding_only_ddG_orchestrator):
# ddG CSV loader does not provide n_chains (it's None)
- assert binding_only_ddG_mc.data.binding.struct_n_chains is None
+ assert binding_only_ddG_orchestrator.data.binding.struct_n_chains is None
@pytest.mark.slow
@@ -447,7 +447,7 @@ def test_binding_only_ddG_svi_runs(tmp_path):
binding_path = _make_binding_only_csv(tmp_path)
ddg_path = _make_ddG_prior_csv(tmp_path)
- mc = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
None, binding_path,
binding_only=True,
theta="thermo.O2_C12_K5_U0_a.PddG",
@@ -455,7 +455,7 @@ def test_binding_only_ddG_svi_runs(tmp_path):
)
out_prefix = str(tmp_path / "svi_binding_only_ddg")
- inference = RunInference(model=mc, seed=0)
+ inference = RunInference(model=orchestrator, seed=0)
svi = inference.setup_svi(adam_step_size=1e-3)
svi_state, params, converged = inference.run_optimization(
svi=svi,
@@ -486,7 +486,7 @@ def test_binding_only_minibatch_get_random_idx_does_not_crash(tmp_path):
binding_path = _make_binding_only_csv(tmp_path, genotypes=big_genotypes)
ddg_path = _make_ddG_prior_csv(tmp_path, genotypes=big_genotypes)
- mc = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
None, binding_path,
binding_only=True,
theta="thermo.O2_C12_K5_U0_a.PddG",
@@ -495,12 +495,12 @@ def test_binding_only_minibatch_get_random_idx_does_not_crash(tmp_path):
)
# Should not raise
- idx = mc.get_random_idx(batch_key=42, num_batches=1)
+ idx = orchestrator.get_random_idx(batch_key=42, num_batches=1)
assert len(idx) == 2
# All genotypes are subsamplable; none are pinned
- assert mc.data.num_binding == 0
- assert len(mc.data.not_binding_idx) == 4
+ assert orchestrator.data.num_binding == 0
+ assert len(orchestrator.data.not_binding_idx) == 4
@pytest.mark.slow
@@ -513,7 +513,7 @@ def test_binding_only_minibatch_svi_runs(tmp_path):
binding_path = _make_binding_only_csv(tmp_path, genotypes=big_genotypes)
ddg_path = _make_ddG_prior_csv(tmp_path, genotypes=big_genotypes)
- mc = ModelOrchestrator(
+ orchestrator = ModelOrchestrator(
None, binding_path,
binding_only=True,
theta="thermo.O2_C12_K5_U0_a.PddG",
@@ -522,7 +522,7 @@ def test_binding_only_minibatch_svi_runs(tmp_path):
)
out_prefix = str(tmp_path / "svi_binding_only_minibatch")
- inference = RunInference(model=mc, seed=0)
+ inference = RunInference(model=orchestrator, seed=0)
svi = inference.setup_svi(adam_step_size=1e-3)
svi_state, params, converged = inference.run_optimization(
svi=svi,
diff --git a/tests/tfscreen/util/validation/test_check.py b/tests/tfscreen/util/validation/test_check.py
index 97b12d19..1f73b171 100644
--- a/tests/tfscreen/util/validation/test_check.py
+++ b/tests/tfscreen/util/validation/test_check.py
@@ -1,5 +1,5 @@
import pytest
-from tfscreen.util.validation.check import check_number
+from tfscreen.util.validation.check import check_number, check_unknown_keys
# ----------------------------------------------------------------------------
# test check_number
@@ -39,4 +39,37 @@ def test_check_number_success(value, kwargs, expected):
def test_check_number_failures(value, kwargs, error, match):
"""Tests that check_number correctly raises errors for invalid inputs."""
with pytest.raises(error, match=match):
- check_number(value, **kwargs)
\ No newline at end of file
+ check_number(value, **kwargs)
+
+
+# ----------------------------------------------------------------------------
+# test check_unknown_keys
+# ----------------------------------------------------------------------------
+
+ALLOWED = frozenset({"alpha", "beta", "gamma"})
+
+
+def test_check_unknown_keys_all_valid():
+ check_unknown_keys({"alpha": 1, "beta": 2}, ALLOWED)
+
+
+def test_check_unknown_keys_empty_config():
+ check_unknown_keys({}, ALLOWED)
+
+
+def test_check_unknown_keys_single_unknown():
+ with pytest.raises(ValueError, match="typo"):
+ check_unknown_keys({"alpha": 1, "typo": 99}, ALLOWED)
+
+
+def test_check_unknown_keys_multiple_unknown():
+ with pytest.raises(ValueError) as exc_info:
+ check_unknown_keys({"alpha": 1, "bad1": 2, "bad2": 3}, ALLOWED)
+ msg = str(exc_info.value)
+ assert "bad1" in msg
+ assert "bad2" in msg
+
+
+def test_check_unknown_keys_label_in_message():
+ with pytest.raises(ValueError, match="my_label"):
+ check_unknown_keys({"unknown": 1}, ALLOWED, label="my_label")
\ No newline at end of file