Bnb/obs model

[bnb32]

• Moved proxy obs sampling logic from Sup3rGanWithObs model into Sampler objects, and removed Sup3rGanWithObs. Sup3rGan standard model now can train with obs with the appropriate batch handler arguments.
• Reworked feature_sets - now uses lr_features, hr_exo_features, and hr_out_features instead of lr_only_features + inferring hr_out_features. This feels simpler to understand for setting up model inputs and outputs.
• Enabled training with just sparse ground truth data. Discrimination will ignore observation features when sparse_disc=False. Disc models could, in theory, accommodate sparse inputs with appropriate layers, like the generator models.
• Removed auto calculation of enhancement factors - doesn't work with UNets or other models with enhancement + downsampling.
• Renamed input_features to gen_features and added true_features in Sup3rLoss class.

[Copilot]

Pull request overview

This PR refactors how sup3r models train and run with observation (obs) features by moving proxy obs sampling into Sampler/DualSampler, removing the dedicated Sup3rGanWithObs model, and standardizing feature configuration via explicit lr_features, hr_exo_features, and hr_out_features.

Changes:

• Move proxy observation masking/appending into sampler objects and remove Sup3rGanWithObs in favor of configuring Sup3rGan via batch handler feature sets/meta.
• Redesign loss feature selection (input_features → gen_features, add true_features) and introduce an obs-masking loss for sparse ground truth.
• Update pipelines/forward-pass/training code and tests to use set_model_params() and the new feature set schema; remove auto enhancement-factor inference.

Reviewed changes

Copilot reviewed 46 out of 47 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/utilities/test_loss_metrics.py	Update loss metric tests for gen_features and set_model_params() usage.
tests/training/test_train_with_obs.py	Add training tests covering conditioned obs and sparse ground truth workflows.
tests/training/test_train_solar.py	Update solar training tests for new feature sets and t_enhance handling.
tests/training/test_train_gan.py	Update GAN training tests for gen_features and set_model_params() calls.
tests/training/test_train_exo_dc.py	Update exo DC tests to explicit lr_features/hr_out_features/hr_exo_features.
tests/training/test_train_exo_cc.py	Update exo CC tests for new feature set semantics and output expectations.
tests/training/test_train_exo.py	Update exo tests for new feature set semantics and output expectations.
tests/training/test_train_conditioned_obs.py	Remove tests tied to deleted Sup3rGanWithObs.
tests/training/test_train_conditional_exo.py	Update conditional exo tests to new feature sets.
tests/samplers/test_with_obs.py	Add sampler tests validating proxy obs masking/append behavior.
tests/samplers/test_feature_sets.py	Update/add feature set validation tests for new schema (incl. obs).
tests/rasterizers/test_dual.py	Update dual rasterizer tests to new .data.low_res/.data.high_res access.
tests/pipeline/test_pipeline.py	Update pipeline tests to use set_model_params() instead of direct meta edits.
tests/pipeline/test_cli.py	Update CLI pipeline tests to use set_model_params() and minor formatting.
tests/output/test_qa.py	Update QA test setup via set_model_params().
tests/forward_pass/test_surface_model.py	Relax enhancement-factor type assertions to accept numpy integer types.
tests/forward_pass/test_multi_step.py	Update multi-step tests to consistently use set_model_params() with enhance factors/resolution.
tests/forward_pass/test_forward_pass_obs.py	Switch forward-pass obs test from Sup3rGanWithObs to Sup3rGan + params.
tests/forward_pass/test_forward_pass_exo.py	Update forward-pass exo tests for set_model_params(); ensure hr_exo_features present in meta in some cases.
tests/forward_pass/test_forward_pass.py	Minor formatting and multi-step model setup via set_model_params().
tests/conftest.py	Rename/add generator configs for obs-aware models used in tests.
tests/bias/test_qdm_bias_correction.py	Formatting and update model meta setup to set_model_params().
tests/bias/test_presrat_bias_correction.py	Update model meta setup to set_model_params().
tests/bias/test_bias_correction.py	Update model meta setup to set_model_params().
tests/batch_queues/test_bq_general.py	Update batch queue test feature set config to new schema.
tests/batch_handlers/test_bh_h5_cc.py	Update CC batch handler tests for new feature set schema.
sup3r/utilities/loss_metrics.py	Introduce Sup3rLoss with gen_features/true_features; add obs-masking loss variant.
sup3r/preprocessing/samplers/dual.py	Refactor DualSampler to support proxy obs and new feature set schema.
sup3r/preprocessing/samplers/dc.py	Update docstrings/expectations to new feature set schema.
sup3r/preprocessing/samplers/cc.py	Update CC sampler docs and shape consistency method naming.
sup3r/preprocessing/samplers/base.py	Add proxy obs generation/append logic and new feature set handling.
sup3r/preprocessing/rasterizers/exo.py	Improve rasterizer docs; route _obs features to ObsRasterizer in factory.
sup3r/preprocessing/rasterizers/dual.py	Refactor to use self.data consistently for dual rasterizer internals.
sup3r/preprocessing/data_handlers/exo.py	Improve exo handler docs and minor validation tweak.
sup3r/preprocessing/collections/stats.py	Add validation when low-res stats requested but containers lack low-res data.
sup3r/preprocessing/batch_queues/dual.py	Simplify dual batch queue shape to LR/HR pair only.
sup3r/preprocessing/batch_queues/base.py	Update single batch queue to use hr_source_features and new LR/HR feature indexing.
sup3r/preprocessing/batch_handlers/factory.py	Update batch handler docs to new feature set schema.
sup3r/models/with_obs.py	Remove Sup3rGanWithObs implementation.
sup3r/models/utilities.py	Consolidate SUP3R layer tuple used for exo/obs layer handling.
sup3r/models/interface.py	Remove auto enhancement inference; add feature consistency check vs layer features; update set_model_params signature/behavior.
sup3r/models/conditional.py	Update training setup to call set_model_params(..., batch_handler=...) (incl. lower models).
sup3r/models/base.py	Add sparse_disc behavior for discriminator input; adjust content loss handling and adversarial combination.
sup3r/models/abstract.py	Update loss feature selection plumbing; adjust public generate() exo layer handling and error reporting.
sup3r/models/init.py	Remove Sup3rGanWithObs export.
pyproject.toml	Bump NREL-rex minimum version.
pixi.lock	Update lockfile for new dependency constraints/version metadata.

Comments suppressed due to low confidence (1)

sup3r/models/abstract.py:1234

• The docstring says obs features at the end of hi_res_true are extracted/trimmed before loss calculation, but the implementation does not do this. If proxy obs channels contain NaNs, standard loss functions (e.g., MAE/MSE) will become NaN when those channels are included in hi_res_true. Either actually split obs channels into hi_res_exo and remove them from hi_res_true before calling calc_loss, or ensure the default content loss only operates on dense hr_out_features.

    def _get_hr_exo_and_loss(
        self,
        low_res,
        hi_res_true,
        **calc_loss_kwargs,
    ):
        """Get high-resolution exogenous data, generate synthetic output, and
        compute loss. Obs features (if present at the end of hi_res_true) are
        extracted and added to exo_data, and trimmed from hi_res_true before
        loss calculation."""
        hi_res_exo = self.get_hr_exo_input(hi_res_true)
        hi_res_gen = self._tf_generate(low_res, hi_res_exo)
        loss, loss_details = self.calc_loss(
            hi_res_true, hi_res_gen, **calc_loss_kwargs
        )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[Copilot]

Pull request overview

Refactors Sup3r observation support by removing the specialized Sup3rGanWithObs model and shifting proxy-observation sampling/masking responsibilities into Sampler objects, while also simplifying/clarifying feature-set configuration across training and forward-pass workflows.

Changes:

• Moved proxy observation generation to Sampler/DualSampler via proxy_obs_kwargs and removed Sup3rGanWithObs.
• Reworked feature_sets to use lr_features, hr_exo_features, and hr_out_features explicitly.
• Updated loss API to use gen_features/true_features, added sparse-obs-capable dummy loss, and updated tests accordingly.

Reviewed changes

Copilot reviewed 47 out of 48 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/utilities/test_loss_metrics.py	Updates loss config args to gen_features and uses set_model_params() in tests.
tests/training/test_train_with_obs.py	Adds training tests for conditioned obs and sparse ground-truth obs workflows.
tests/training/test_train_solar.py	Updates solar CC tests for new feature_sets and passes t_enhance.
tests/training/test_train_gan.py	Updates custom loss config and ensures set_model_params() is used.
tests/training/test_train_exo_dc.py	Updates DC exogenous feature_sets to new schema.
tests/training/test_train_exo_cc.py	Updates CC exogenous tests to new feature_sets and output assertions.
tests/training/test_train_exo.py	Updates non-CC exogenous tests to new feature_sets and output assertions.
tests/training/test_train_conditioned_obs.py	Removes legacy Sup3rGanWithObs conditioned-obs training test.
tests/training/test_train_conditional_exo.py	Updates conditional-exo test feature_sets to new schema.
tests/samplers/test_with_obs.py	Adds tests validating proxy-obs sampling behavior and fractions.
tests/samplers/test_feature_sets.py	Updates feature-set validation tests for new schema and obs inclusion.
tests/rasterizers/test_dual.py	Updates dual rasterizer attribute access via data.low_res/high_res.
tests/pipeline/test_pipeline.py	Uses set_model_params() instead of directly mutating meta.
tests/pipeline/test_cli.py	Uses set_model_params() and minor formatting cleanup.
tests/output/test_qa.py	Uses set_model_params() for model metadata setup.
tests/forward_pass/test_surface_model.py	Allows numpy integer types for enhancement-factor assertions.
tests/forward_pass/test_multi_step.py	Uses set_model_params() with explicit enhance factors/resolution.
tests/forward_pass/test_forward_pass_obs.py	Switches from Sup3rGanWithObs to Sup3rGan + hr_exo_features.
tests/forward_pass/test_forward_pass_exo.py	Uses set_model_params() in several forward-pass integration tests; adds hr_exo_features meta where needed.
tests/forward_pass/test_forward_pass.py	Minor formatting tweaks and uses set_model_params() in multi-step setup.
tests/conftest.py	Renames/introduces generator configs that include obs layers.
tests/bias/test_qdm_bias_correction.py	Formatting + uses set_model_params() in FWP integration test.
tests/bias/test_presrat_bias_correction.py	Uses set_model_params() in FWP integration test.
tests/bias/test_bias_correction.py	Uses set_model_params() in FWP integration test.
tests/batch_queues/test_bq_general.py	Updates batch-queue test feature_sets to new schema.
tests/batch_handlers/test_bh_h5_cc.py	Updates CC batch-handler tests to new feature_sets schema.
sup3r/utilities/loss_metrics.py	Introduces Sup3rLoss base and adds gen_features/true_features support; adds sparse-obs dummy loss.
sup3r/preprocessing/samplers/dual.py	Adds proxy_obs_kwargs, updates feature-set handling, and refactors init/consistency checks.
sup3r/preprocessing/samplers/dc.py	Updates docstrings for new feature_sets schema.
sup3r/preprocessing/samplers/cc.py	Updates docstrings and shape-consistency method naming.
sup3r/preprocessing/samplers/base.py	Adds proxy-obs generation/masking, new feature-set semantics, and related consistency checks.
sup3r/preprocessing/rasterizers/exo.py	Improves docstrings and clarifies rasterizer selection behavior.
sup3r/preprocessing/rasterizers/dual.py	Refactors to store dataset under self.data and updates references.
sup3r/preprocessing/data_handlers/exo.py	Docstring improvements and minor assertion style update.
sup3r/preprocessing/collections/stats.py	Adds validation when low-res features are requested but containers lack low-res data.
sup3r/preprocessing/batch_queues/dual.py	Simplifies queue shape to align with two-member dual batches.
sup3r/preprocessing/batch_queues/base.py	Updates single-queue shape and transform logic to use new feature indices.
sup3r/preprocessing/batch_handlers/factory.py	Updates docs for new feature_sets schema.
sup3r/models/with_obs.py	Removes Sup3rGanWithObs implementation.
sup3r/models/utilities.py	Consolidates SUP3R_LAYERS tuple definition.
sup3r/models/solar_cc.py	Propagates train_disc into init_weights() and minor signature/style updates.
sup3r/models/interface.py	Removes auto-enhance computation, adds exo-feature consistency check, updates set_model_params() behavior.
sup3r/models/conditional.py	Uses set_model_params(batch_handler=...) for self and lower models; minor logging string fix.
sup3r/models/base.py	Skips discriminator init/compute when disabled; refactors loss computation paths.
sup3r/models/abstract.py	Updates loss feature-selection plumbing and exogenous feature handling/validation.
sup3r/models/init.py	Removes export of Sup3rGanWithObs.
pyproject.toml	Bumps NREL-rex minimum version.
pixi.lock	Updates lockfile entries to reflect dependency/version changes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.