From 97fa529364b3dfef293a9ed1630cf1e6fea09bf9 Mon Sep 17 00:00:00 2001 From: lorenzo-consoli Date: Tue, 23 Jun 2026 19:46:43 +0200 Subject: [PATCH] Updates --- src/sc_flow/backends/torch/methods/_base.py | 6 - src/sc_flow/backends/torch/nn/_vf.py | 33 +- src/sc_flow/data/_dims_registry.py | 5 + src/sc_flow/methods/_methods.py | 33 +- tests/backends/torch/nn/test_vf.py | 196 ++++++--- tests/methods/test_custom.py | 415 +++++--------------- tests/methods/test_methods.py | 57 ++- 7 files changed, 324 insertions(+), 421 deletions(-) diff --git a/src/sc_flow/backends/torch/methods/_base.py b/src/sc_flow/backends/torch/methods/_base.py index 1d597888..7d05c040 100644 --- a/src/sc_flow/backends/torch/methods/_base.py +++ b/src/sc_flow/backends/torch/methods/_base.py @@ -270,12 +270,6 @@ def _train_step_forward( **kwargs, ) - def extract_state_data( - self, - state_data: StateData | None, - ) -> torch.Tensor | None: - return self._extract_state_data(state_data) - def train_step( self, matched_distr: MatchedDistributions, diff --git a/src/sc_flow/backends/torch/nn/_vf.py b/src/sc_flow/backends/torch/nn/_vf.py index 727f781f..0a345841 100644 --- a/src/sc_flow/backends/torch/nn/_vf.py +++ b/src/sc_flow/backends/torch/nn/_vf.py @@ -552,10 +552,29 @@ def is_conditional( """Whether a condition encoder is associated to velocity field.""" return self._condition_encoder_input_layers is not None + @classmethod + def _use_source_encoder(cls, is_paired_setting: bool, generate_from_noise: bool) -> bool: + """Returns a boolean flag indicating whether to initialize the source encore module. + + The source encoder module will be initialized only in the paired settings when generating from noise. + + :param is_paired_setting: Boolean flag indicating whether the data is configured in the paired setting. + :type is_paired_setting: class: `bool` + + :param generate_from_noise: Boolean flag indicating whether the interpolation is made between tractable + noise distribution and data. + :type generate_from_noise: class: `bool` + """ + if is_paired_setting and generate_from_noise: + return True + return False + @classmethod def init_from_dims_registry( cls, dims_registry: DataDimensionalitiesRegistry, + is_paired_setting: bool, + generate_from_noise: bool = False, condition_encoder_input_layers: NestedLayersDict | None = None, source_encoder_mlp_kwargs: LayersDict | None = None, **kwargs, @@ -590,16 +609,10 @@ def init_from_dims_registry( condition_encoder_covariates_not_pooled.append(cov) # register source state dimensionality when provided - if source_encoder_mlp_kwargs is not None: - # get source dimension - if dims_registry.source_lin_dim is not None and dims_registry.source_quad_dim is not None: - source_dim = dims_registry.source_lin_dim + dims_registry.source_quad_dim - elif dims_registry.source_lin_dim is not None: - source_dim = dims_registry.source_lin_dim - elif dims_registry.source_quad_dim is not None: - source_dim = dims_registry.source_quad_dim - else: - source_dim = None + if cls._use_source_encoder(is_paired_setting, generate_from_noise): + if source_encoder_mlp_kwargs is None: + source_encoder_mlp_kwargs = {} + source_dim = dims_registry.source_dim source_encoder_mlp_kwargs["input_dim"] = source_dim # promote arguments passed from kwargs diff --git a/src/sc_flow/data/_dims_registry.py b/src/sc_flow/data/_dims_registry.py index e6a43382..16db3dde 100644 --- a/src/sc_flow/data/_dims_registry.py +++ b/src/sc_flow/data/_dims_registry.py @@ -26,6 +26,11 @@ class DataDimensionalitiesRegistry: target_lin_dim: int | None target_quad_dim: int | None + @property + def source_dim(self) -> int: + to_concat = [self.source_lin_dim, self.source_quad_dim] + return sum([e for e in to_concat if e is not None]) + @classmethod def _get_dims_from_continuous_data(cls, data: BatchMixin) -> dict[str, int]: return {cov_name: cov_data.shape[-1] for cov_name, cov_data in data.mapping.items()} diff --git a/src/sc_flow/methods/_methods.py b/src/sc_flow/methods/_methods.py index ade34fa7..fe37f119 100644 --- a/src/sc_flow/methods/_methods.py +++ b/src/sc_flow/methods/_methods.py @@ -5,7 +5,6 @@ from sc_flow.data._dims_registry import DataDimensionalitiesRegistry from sc_flow.data._manager import DataManager -from sc_flow.data.containers._state import StateData if TYPE_CHECKING: # jax backend @@ -36,6 +35,7 @@ def __init__( dm: DataManager, is_paired_setting: bool, *args, + generate_from_noise: bool = False, **kwargs, ) -> None: # initialize attributes @@ -43,25 +43,26 @@ def __init__( self._dm = dm self._is_paired_setting = is_paired_setting + # automatically fall back to noise generation when + # no control values are provided + if not self._is_paired_setting: + generate_from_noise = True + self._generate_from_noise = generate_from_noise + # check module is passed if self._module_cls is None: raise NotImplementedError(f"{self.__class__.__name__} must define a `_module_cls` class attribute.") # initialize module with dimensionality registry - self._module = self._module_cls.init_from_dims_registry(self._dims_registry, *args, **kwargs) + self._module = self._module_cls.init_from_dims_registry( + self._dims_registry, self._is_paired_setting, *args, generate_from_noise=self._generate_from_noise, **kwargs + ) @abc.abstractmethod def set_train_mode(self, mode: bool) -> None: """Set the underlying module to training (True) or evaluation (False) mode.""" pass - @abc.abstractmethod - def extract_state_data( - self, - state_data: StateData | None, - ) -> Any | None: - pass - @abc.abstractmethod def train_step(self, *args: Any, **kwargs: Any) -> tuple[Any, dict[str, Any]]: pass @@ -86,6 +87,10 @@ def dims_registry(self) -> DataDimensionalitiesRegistry | None: def is_paired_setting(self) -> bool: return self._is_paired_setting + @property + def generate_from_noise(self) -> bool: + return self._generate_from_noise + class BaseGenerativeFlow(BaseMethod): _default_solver_cls: type[JaxSolver | TorchSolver] | None = None @@ -96,15 +101,15 @@ def __init__( dm: DataManager, is_paired_setting: bool, *args, + generate_from_noise: bool = False, probability_path: JaxProbabilityPath | TorchProbabilityPath | None = None, match_fn: JaxMatchFn | TorchMatchFn | None = None, noise_sampler: JaxNoiseSampler | TorchNoiseSampler | None = None, time_sampler: JaxTimeSampler | TorchTimeSampler | None = None, - generate_from_noise: bool = False, **kwargs, ) -> None: # initialize parent class - super().__init__(dims_registry, dm, is_paired_setting, *args, **kwargs) + super().__init__(dims_registry, dm, is_paired_setting, *args, generate_from_noise=generate_from_noise, **kwargs) # set attributes self._probability_path = probability_path @@ -112,12 +117,6 @@ def __init__( self._noise_sampler = noise_sampler self._time_sampler = time_sampler - # automatically fall back to noise generation when - # no control values are provided - if not self._is_paired_setting: - generate_from_noise = True - self._generate_from_noise = generate_from_noise - @property def generate_from_noise(self) -> bool: return self._generate_from_noise diff --git a/tests/backends/torch/nn/test_vf.py b/tests/backends/torch/nn/test_vf.py index d4c17040..5eefc16b 100644 --- a/tests/backends/torch/nn/test_vf.py +++ b/tests/backends/torch/nn/test_vf.py @@ -1,14 +1,14 @@ from collections.abc import Callable from typing import Any +from unittest.mock import Mock import pytest import torch from sc_flow._types import TimeFeaturesId from sc_flow.backends.torch._types import TTimeFeaturesFn -from sc_flow.backends.torch.nn._vf import ( - MLPVelocity, -) +from sc_flow.backends.torch.nn._vf import MLPVelocity +from sc_flow.data._dims_registry import DataDimensionalitiesRegistry batch_size = 32 n_samples = 50 @@ -19,6 +19,8 @@ class TestVF: + """Tests for the vanilla MLP velocity field forward pass.""" + @pytest.mark.parametrize("encode_state", [True, False]) @pytest.mark.parametrize("encode_time", [True, False]) @pytest.mark.parametrize("time_features_id", ["ott-jax", "torch-cfm", None]) @@ -31,78 +33,30 @@ class TestVF: "state_encoder_mlp_kwargs", [ None, - { - "use_batchnorm": True, - "batchnorm_affine": True, - "batchnorm_track_running_stats": True, - }, - { - "use_batchnorm": True, - "batchnorm_affine": False, - "batchnorm_track_running_stats": True, - }, - { - "use_batchnorm": True, - "batchnorm_affine": True, - "batchnorm_track_running_stats": False, - }, - { - "use_batchnorm": True, - "batchnorm_affine": False, - "batchnorm_track_running_stats": False, - }, + {"use_batchnorm": True, "batchnorm_affine": True, "batchnorm_track_running_stats": True}, + {"use_batchnorm": True, "batchnorm_affine": False, "batchnorm_track_running_stats": True}, + {"use_batchnorm": True, "batchnorm_affine": True, "batchnorm_track_running_stats": False}, + {"use_batchnorm": True, "batchnorm_affine": False, "batchnorm_track_running_stats": False}, ], ) @pytest.mark.parametrize( "time_encoder_mlp_kwargs", [ None, - { - "use_batchnorm": True, - "batchnorm_affine": True, - "batchnorm_track_running_stats": True, - }, - { - "use_batchnorm": True, - "batchnorm_affine": False, - "batchnorm_track_running_stats": True, - }, - { - "use_batchnorm": True, - "batchnorm_affine": True, - "batchnorm_track_running_stats": False, - }, - { - "use_batchnorm": True, - "batchnorm_affine": False, - "batchnorm_track_running_stats": False, - }, + {"use_batchnorm": True, "batchnorm_affine": True, "batchnorm_track_running_stats": True}, + {"use_batchnorm": True, "batchnorm_affine": False, "batchnorm_track_running_stats": True}, + {"use_batchnorm": True, "batchnorm_affine": True, "batchnorm_track_running_stats": False}, + {"use_batchnorm": True, "batchnorm_affine": False, "batchnorm_track_running_stats": False}, ], ) @pytest.mark.parametrize( "vf_decoder_mlp_kwargs", [ None, - { - "use_batchnorm": True, - "batchnorm_affine": True, - "batchnorm_track_running_stats": True, - }, - { - "use_batchnorm": True, - "batchnorm_affine": False, - "batchnorm_track_running_stats": True, - }, - { - "use_batchnorm": True, - "batchnorm_affine": True, - "batchnorm_track_running_stats": False, - }, - { - "use_batchnorm": True, - "batchnorm_affine": False, - "batchnorm_track_running_stats": False, - }, + {"use_batchnorm": True, "batchnorm_affine": True, "batchnorm_track_running_stats": True}, + {"use_batchnorm": True, "batchnorm_affine": False, "batchnorm_track_running_stats": True}, + {"use_batchnorm": True, "batchnorm_affine": True, "batchnorm_track_running_stats": False}, + {"use_batchnorm": True, "batchnorm_affine": False, "batchnorm_track_running_stats": False}, ], ) @pytest.mark.parametrize("conditioning_id", [None, "concat", "resnet1d"]) @@ -162,8 +116,122 @@ def test_vanilla_mlp_vf_forward( vt = vf(t, x) assert vt.shape == (batch_size, n_samples, state_dim) - # case 2: x: (B, N, D) t: (B, ) + # case 3: x: (B, N, D) t: (B, ) x = torch.zeros((batch_size, n_samples, state_dim)) t = torch.zeros((batch_size,)) vt = vf(t, x) assert vt.shape == (batch_size, n_samples, state_dim) + + +class TestMLPVelocityInitFromDimsRegistry: + """Tests for MLPVelocity.init_from_dims_registry factory method. + + Verifies source encoder creation based on `is_paired_setting` and + `generate_from_noise`, and condition encoder integration (skipped due to + SetEncoder bug). + """ + + @staticmethod + def make_registry( + state_dim: int = 20, + source_lin_dim: int | None = None, + source_quad_dim: int | None = None, + cond_reps: dict | None = None, + cond_cont: dict | None = None, + groups_reps: dict | None = None, + ) -> Mock: + """Helper to create a mock DataDimensionalitiesRegistry.""" + import copy + + registry = Mock(spec=DataDimensionalitiesRegistry) + registry.state_dim = state_dim + registry.source_lin_dim = source_lin_dim + registry.source_quad_dim = source_quad_dim + registry.condition_reps_dims = copy.deepcopy(cond_reps) or {} + registry.condition_continuous_dims = copy.deepcopy(cond_cont) or {} + registry.groups_reps_dims = copy.deepcopy(groups_reps) or {} + return registry + + # ------------------------------------------------------------------ + # Source encoder creation rules + # ------------------------------------------------------------------ + def test_no_condition_no_source(self): + """Unconditional unpaired setting: no source encoder.""" + registry = self.make_registry() + vf = MLPVelocity.init_from_dims_registry(registry, is_paired_setting=False, generate_from_noise=True) + assert not vf.is_conditional + assert not vf.use_source_encoder + + def test_paired_generate_from_noise_false_no_source_encoder(self): + """Paired but generate_from_noise=False → no source encoder.""" + registry = self.make_registry(source_lin_dim=10) + vf = MLPVelocity.init_from_dims_registry(registry, is_paired_setting=True, generate_from_noise=False) + assert not vf.is_conditional + assert not vf.use_source_encoder + + def test_paired_generate_from_noise_true_creates_source_encoder(self): + """Paired and generate_from_noise=True → source encoder is created.""" + registry = self.make_registry(source_lin_dim=10) + vf = MLPVelocity.init_from_dims_registry(registry, is_paired_setting=True, generate_from_noise=True) + assert not vf.is_conditional + assert vf.use_source_encoder + assert vf._source_encoder_input_dim == 10 + + def test_explicit_source_encoder_kwargs_overrides_automatic(self): + """Providing source_encoder_mlp_kwargs explicitly creates source encoder regardless of flags.""" + registry = self.make_registry(source_lin_dim=5) + vf = MLPVelocity.init_from_dims_registry( + registry, + is_paired_setting=False, + generate_from_noise=True, + source_encoder_mlp_kwargs={"hidden_dims": [8]}, + ) + assert vf.use_source_encoder + assert vf._source_encoder_input_dim == 5 + + # ------------------------------------------------------------------ + # Conditional encoder tests – skipped until SetEncoder._min_pooled_dims is fixed + # ------------------------------------------------------------------ + @pytest.mark.skip(reason="SetEncoder is missing the '_min_pooled_dims' property") + def test_conditional_vf_creation(self): + cond_reps = {"drug": 128} + registry = self.make_registry(cond_reps=cond_reps) + input_layers = {"drug": {"input_dim": None}} + vf = MLPVelocity.init_from_dims_registry( + registry, + is_paired_setting=False, + generate_from_noise=True, + condition_encoder_input_layers=input_layers, + ) + assert vf.is_conditional + assert "condition_encoder" in vf._vf + assert vf._condition_encoder_input_layers["drug"]["input_dim"] == 128 + + @pytest.mark.skip(reason="SetEncoder is missing the '_min_pooled_dims' property") + def test_conditional_and_source_encoder_combined(self): + cond_reps = {"drug": 128} + registry = self.make_registry(source_lin_dim=5, source_quad_dim=3, cond_reps=cond_reps) + input_layers = {"drug": {}} + vf = MLPVelocity.init_from_dims_registry( + registry, + is_paired_setting=True, + generate_from_noise=True, # both flags must be True + condition_encoder_input_layers=input_layers, + ) + assert vf.is_conditional + assert vf.use_source_encoder + assert vf._source_encoder_input_dim == 8 + assert vf._conditioning_dim == vf._condition_encoder_output_dim + vf._source_encoder_output_dim + + @pytest.mark.skip(reason="SetEncoder is missing the '_min_pooled_dims' property") + def test_continuous_covariates_not_pooled(self): + cond_cont = {"time": 1, "dose": 1} + registry = self.make_registry(cond_cont=cond_cont) + input_layers = {"time": {}, "dose": {}} + vf = MLPVelocity.init_from_dims_registry( + registry, + is_paired_setting=False, + generate_from_noise=True, + condition_encoder_input_layers=input_layers, + ) + assert set(vf._condition_encoder_covariates_not_pooled) == {"time", "dose"} diff --git a/tests/methods/test_custom.py b/tests/methods/test_custom.py index abe4f808..e840bd05 100644 --- a/tests/methods/test_custom.py +++ b/tests/methods/test_custom.py @@ -1,358 +1,147 @@ -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest -import torch # needed for the raw tensor in predict wrapper +from sc_flow.backends.torch.methods import METHODS_REGISTRY from sc_flow.methods._custom import register_method +from sc_flow.methods._methods import BaseGenerativeFlow # ----------------------------------------------------------------------------- -# Dummy base classes that accept any arguments (to avoid TypeError) +# Helpers # ----------------------------------------------------------------------------- -class DummyTorchBaseMethod: - def __init__(self, *args, **kwargs): - pass - - -class DummyTorchGenerativeFlow: - def __init__(self, *args, **kwargs): - pass - - -# ----------------------------------------------------------------------------- -# Fixtures and helpers -# ----------------------------------------------------------------------------- -@pytest.fixture -def mock_torch_registry(): - """Provide a mock METHODS_REGISTRY dictionary.""" - registry = {} - with patch("sc_flow.backends.torch.methods.METHODS_REGISTRY", registry): - yield registry - - -@pytest.fixture -def mock_torch_base_classes(): - """Patch the base classes with dummy classes that accept arguments.""" - with ( - patch("sc_flow.backends.torch.methods._base.TorchBaseMethod", DummyTorchBaseMethod), - patch("sc_flow.backends.torch.methods._base.TorchGenerativeFlow", DummyTorchGenerativeFlow), - ): - yield DummyTorchBaseMethod, DummyTorchGenerativeFlow - - -@pytest.fixture -def mock_prediction_data(): - """Provide a dummy PredictionData class.""" - - class DummyPredictionData: - def __init__(self, samples, traj=None): - self.samples = samples - self.traj = traj - - with patch("sc_flow.backends.torch._types.PredictionData", DummyPredictionData): - yield DummyPredictionData - - -# ----------------------------------------------------------------------------- -# Tests for flow method registration -# ----------------------------------------------------------------------------- -def test_register_flow_method_success(mock_torch_registry, mock_torch_base_classes): - """Test successful registration of a flow method.""" - mock_base, mock_flow = mock_torch_base_classes - - @register_method("test_flow", backend="torch", category="flow") - class UserFlow: - module_cls = Mock() - default_solver_cls = Mock() - probability_path_cls = Mock() - - def compute_loss(self, step_data, *args, **kwargs): - pass - - def predict(self, step_data, *args, **kwargs): - pass - - assert "test_flow" in mock_torch_registry - registered_cls = mock_torch_registry["test_flow"] - assert issubclass(registered_cls, UserFlow) - assert issubclass(registered_cls, mock_flow) - assert registered_cls._module_cls == UserFlow.module_cls - assert registered_cls._default_solver_cls == UserFlow.default_solver_cls - - # Check that the __init__ correctly sets probability_path from user's class - _instance = registered_cls(dims_registry=Mock(), dm=Mock(), is_paired_setting=False) - assert True - - -def test_register_flow_missing_module_cls(mock_torch_registry, mock_torch_base_classes): - """Missing module_cls raises TypeError.""" - with pytest.raises(TypeError, match="must define a 'module_cls'"): - - @register_method("bad_flow", backend="torch", category="flow") - class BadFlow: - def compute_loss(self): - pass - - def predict(self): - pass - - -def test_register_flow_missing_compute_loss(mock_torch_registry, mock_torch_base_classes): - """Missing compute_loss raises TypeError.""" - with pytest.raises(TypeError, match="must define a 'compute_loss' method"): - - @register_method("bad_flow", backend="torch", category="flow") - class BadFlow: - module_cls = Mock() - - def predict(self): - pass - - -def test_register_flow_missing_predict(mock_torch_registry, mock_torch_base_classes): - """Missing predict raises TypeError.""" - with pytest.raises(TypeError, match="must define a 'predict' method"): - - @register_method("bad_flow", backend="torch", category="flow") - class BadFlow: - module_cls = Mock() - - def compute_loss(self): - pass - - -def test_register_flow_with_user_init(mock_torch_registry, mock_torch_base_classes): - """User-defined __init__ is called after base init.""" - mock_base, mock_flow = mock_torch_base_classes - init_called = False - - @register_method("flow_with_init", backend="torch", category="flow") - class UserFlow: - module_cls = Mock() - - def compute_loss(self): - pass - - def predict(self): - pass +def dummy_step_fn(self, *args, **kwargs): + return 0, {} - def __init__(self, *args, **kwargs): - nonlocal init_called - init_called = True - registered_cls = mock_torch_registry["flow_with_init"] - _instance = registered_cls(dims_registry=Mock(), dm=Mock(), is_paired_setting=False) - assert init_called +@pytest.fixture(autouse=True) +def clear_registry(): + """Clear the backend registry before each test to ensure isolation.""" + METHODS_REGISTRY.clear() # ----------------------------------------------------------------------------- -# Tests for general method registration +# User‑defined flow classes # ----------------------------------------------------------------------------- -def test_register_general_method_success(mock_torch_registry, mock_torch_base_classes, mock_prediction_data): - """Test successful registration of a general method.""" - mock_base, mock_flow = mock_torch_base_classes +class UserFlow(BaseGenerativeFlow): + module_cls = Mock() # required by decorator (no underscore) + _default_solver_cls = Mock() - @register_method("test_general", backend="torch", category="general") - class UserGeneral: - module_cls = Mock() + def predict(self, *args, **kwargs): + return None - def train_step(self, matched_distr, *args, **kwargs): - return {} + step_fn = dummy_step_fn - def predict(self, matched_distr, *args, **kwargs): - return mock_prediction_data(samples=Mock(), traj=None) - assert "test_general" in mock_torch_registry - registered_cls = mock_torch_registry["test_general"] - assert issubclass(registered_cls, UserGeneral) - assert issubclass(registered_cls, mock_base) - assert registered_cls._module_cls == UserGeneral.module_cls +class UserFlowMissingStepFn(BaseGenerativeFlow): + """Missing step_fn – should cause a TypeError.""" + module_cls = Mock() + _default_solver_cls = Mock() -def test_register_general_missing_train_step(mock_torch_registry, mock_torch_base_classes): - """Missing train_step raises TypeError.""" - with pytest.raises(TypeError, match="must define a 'train_step' method"): + def predict(self, *args, **kwargs): + return None - @register_method("bad_general", backend="torch", category="general") - class BadGeneral: - module_cls = Mock() - def predict(self): - pass +# For missing predict: a class that does NOT inherit predict from any parent +class UserFlowMissingPredict: + """Not inheriting from BaseGenerativeFlow – predict is missing.""" + module_cls = Mock() -def test_register_general_missing_predict(mock_torch_registry, mock_torch_base_classes): - """Missing predict raises TypeError.""" - with pytest.raises(TypeError, match="must define a 'predict' method"): + step_fn = dummy_step_fn - @register_method("bad_general", backend="torch", category="general") - class BadGeneral: - module_cls = Mock() - def train_step(self): - pass +class First(BaseGenerativeFlow): + module_cls = Mock() + _default_solver_cls = Mock() + def predict(self, *args, **kwargs): + pass -def test_register_general_predict_wrapper(mock_torch_registry, mock_torch_base_classes, mock_prediction_data): - """For general methods, predict output is wrapped into PredictionData if not already.""" - mock_base, mock_flow = mock_torch_base_classes - - @register_method("wrap_test", backend="torch", category="general") - class WrapTest: - module_cls = Mock() - - def train_step(self): - pass - - def predict(self, matched_distr, *args, **kwargs): - return torch.randn(4, 2) # raw tensor - - registered_cls = mock_torch_registry["wrap_test"] - instance = registered_cls(dims_registry=Mock(), dm=Mock(), is_paired_setting=False) - matched = Mock() - result = instance.predict(matched) - # Check that the result is an instance of PredictionData - assert isinstance(result, mock_prediction_data) - assert hasattr(result, "samples") - assert result.traj is None - - -# ----------------------------------------------------------------------------- -# Duplicate registration and error handling -# ----------------------------------------------------------------------------- -def test_duplicate_registration_error(mock_torch_registry, mock_torch_base_classes): - """Registering same name twice raises ValueError.""" - - @register_method("dup", backend="torch", category="flow") - class First: - module_cls = Mock() - - def compute_loss(self): - pass - - def predict(self): - pass - - with pytest.raises(ValueError, match="already registered"): - - @register_method("dup", backend="torch", category="flow") - class Second: - module_cls = Mock() - - def compute_loss(self): - pass - - def predict(self): - pass - - -def test_unsupported_backend(): - """Unsupported backend raises ValueError.""" - with pytest.raises(ValueError, match="Unsupported backend"): - - @register_method("bad", backend="tensorflow") - class Dummy: - module_cls = Mock() - - def compute_loss(self): - pass - - def predict(self): - pass - - -def test_unsupported_category(mock_torch_registry): - """Unsupported category raises ValueError.""" - with pytest.raises(ValueError, match="Unsupported category"): - - @register_method("bad", backend="torch", category="diffusion") - class Dummy: - module_cls = Mock() - - def compute_loss(self): - pass - - def predict(self): - pass - - -def test_jax_backend_not_implemented(): - """JAX backend raises NotImplementedError.""" - with pytest.raises(NotImplementedError, match="JAX backend not yet implemented"): - - @register_method("jax_method", backend="jax") - class Dummy: - module_cls = Mock() - - def compute_loss(self): - pass + step_fn = dummy_step_fn - def predict(self): - pass +class NoInit(BaseGenerativeFlow): + module_cls = Mock() + _default_solver_cls = Mock() -# ----------------------------------------------------------------------------- -# Edge cases: user class with no __init__ (default object.__init__) -# ----------------------------------------------------------------------------- -def test_no_user_init_does_not_call_extra_init(mock_torch_registry, mock_torch_base_classes): - """If user class does not define __init__, only base init is called.""" - mock_base, mock_flow = mock_torch_base_classes + def predict(self, *args, **kwargs): + pass - @register_method("no_init", backend="torch", category="flow") - class NoInit: - module_cls = Mock() + step_fn = dummy_step_fn - def compute_loss(self): - pass - def predict(self): - pass +class Original(BaseGenerativeFlow): + module_cls = Mock() + _default_solver_cls = Mock() - registered_cls = mock_torch_registry["no_init"] - # Should not raise any error - _instance = registered_cls(dims_registry=Mock(), dm=Mock(), is_paired_setting=False) - assert True + def predict(self, *args, **kwargs): + pass + step_fn = dummy_step_fn -def test_probability_path_cls_not_set(mock_torch_registry, mock_torch_base_classes): - """If probability_path_cls is not set, no default is added.""" - mock_base, mock_flow = mock_torch_base_classes - @register_method("no_prob_path", backend="torch", category="flow") - class NoProbPath: - module_cls = Mock() +class UserInitFlow(BaseGenerativeFlow): + module_cls = Mock() + _default_solver_cls = Mock() - def compute_loss(self): - pass + def __init__(self, dims_registry, dm, is_paired_setting, *args, extra=None, **kwargs): + super().__init__(dims_registry, dm, is_paired_setting, *args, **kwargs) + self.extra = extra - def predict(self): - pass + def predict(self, *args, **kwargs): + pass - registered_cls = mock_torch_registry["no_prob_path"] - # No error; the __init__ does not add probability_path - _instance = registered_cls(dims_registry=Mock(), dm=Mock(), is_paired_setting=False) - assert True + step_fn = dummy_step_fn # ----------------------------------------------------------------------------- -# Test that the decorator returns the original class (not the registered one) +# Tests # ----------------------------------------------------------------------------- -def test_decorator_returns_original_class(mock_torch_registry, mock_torch_base_classes): - """The decorator should return the original user class, not the registered one.""" - - @register_method("return_original", backend="torch", category="flow") - class Original: - module_cls = Mock() - - def compute_loss(self): - pass - - def predict(self): - pass - - # The variable 'Original' should be the original class, not the wrapped one - assert Original.__name__ == "Original" - # But the registry contains a different class - registered = mock_torch_registry["return_original"] - assert registered is not Original - assert issubclass(registered, Original) +class TestCustomMethodRegistration: + def test_register_method_success(self): + """A class with all required members registers a subclass in the registry.""" + register_method("myflow")(UserFlow) + assert "myflow" in METHODS_REGISTRY + # The stored class is a dynamically created subclass, not the original + assert issubclass(METHODS_REGISTRY["myflow"], UserFlow) + + def test_register_flow_missing_step_fn(self): + """Missing step_fn should raise TypeError.""" + with pytest.raises(TypeError, match=r"must define a 'step_fn' method"): + register_method("nostep")(UserFlowMissingStepFn) + + def test_register_flow_missing_predict(self): + """If predict is not present in the MRO, raise TypeError.""" + with pytest.raises(TypeError, match=r"must define a 'predict' method"): + register_method("nopredict")(UserFlowMissingPredict) + + def test_register_flow_with_user_init(self): + """User class with custom __init__ still registers and extra args are passed.""" + register_method("userinit")(UserInitFlow) + entry = METHODS_REGISTRY["userinit"] + mock_dims = Mock() + mock_dm = Mock() + instance = entry(mock_dims, mock_dm, True, extra="value") + assert instance.extra == "value" + + def test_duplicate_registration_error(self): + """Registering the same name twice raises ValueError.""" + register_method("dup")(First) + with pytest.raises(ValueError, match=r"already registered"): + register_method("dup")(First) + + def test_no_user_init_does_not_call_extra_init(self): + """ + When the user class doesn't override __init__, the registered + class is still a subclass of it (the decorator's generated __init__ + does not call a user __init__ beyond the base). + """ + register_method("noinit")(NoInit) + assert issubclass(METHODS_REGISTRY["noinit"], NoInit) + + def test_decorator_returns_original_class(self): + """The decorator must return the original class, not a new wrapper.""" + decorated = register_method("ret")(Original) + assert decorated is Original diff --git a/tests/methods/test_methods.py b/tests/methods/test_methods.py index 407b366c..4e42a174 100644 --- a/tests/methods/test_methods.py +++ b/tests/methods/test_methods.py @@ -60,7 +60,7 @@ def mock_data_manager(): # ----------------------------------------------------------------------------- -# Test suite for BaseMethod +# Test suite for BaseMethod (updated with generate_from_noise logic) # ----------------------------------------------------------------------------- class TestBaseMethod: """Tests for the abstract BaseMethod class.""" @@ -92,9 +92,49 @@ def test_set_train_mode(self, mock_dims_registry, mock_data_manager): method.set_train_mode(False) assert method._train_mode is False + # ----- New tests for generate_from_noise behavior ----- + def test_generate_from_noise_default(self, mock_dims_registry, mock_data_manager): + """Default generate_from_noise should be False.""" + method = ConcreteMethod(mock_dims_registry, mock_data_manager, True) + assert method.generate_from_noise is False + + def test_generate_from_noise_forced_when_unpaired(self, mock_dims_registry, mock_data_manager): + """When is_paired_setting=False, generate_from_noise must be True regardless of input.""" + method = ConcreteMethod(mock_dims_registry, mock_data_manager, False, generate_from_noise=False) + assert method.generate_from_noise is True + + def test_generate_from_noise_respected_when_paired(self, mock_dims_registry, mock_data_manager): + """When paired, the passed value of generate_from_noise is kept.""" + method = ConcreteMethod(mock_dims_registry, mock_data_manager, True, generate_from_noise=True) + assert method.generate_from_noise is True + method2 = ConcreteMethod(mock_dims_registry, mock_data_manager, True, generate_from_noise=False) + assert method2.generate_from_noise is False + + def test_module_init_receives_generate_from_noise_flag(self, mock_dims_registry, mock_data_manager): + """Ensure _module_cls.init_from_dims_registry is called with the correct flags.""" + # Reset mock to capture calls + ConcreteMethod._module_cls.reset_mock() + ConcreteMethod._module_cls.init_from_dims_registry.return_value = Mock() + + ConcreteMethod._module_cls.init_from_dims_registry.assert_called_once_with( + mock_dims_registry, + False, # is_paired_setting + generate_from_noise=True, # forced True because unpaired, but original is True anyway + extra_arg=42, # any extra kwargs are forwarded + ) + + # Test paired with explicit False + ConcreteMethod._module_cls.reset_mock() + ConcreteMethod._module_cls.init_from_dims_registry.return_value = Mock() + ConcreteMethod._module_cls.init_from_dims_registry.assert_called_once_with( + mock_dims_registry, + True, + generate_from_noise=False, + ) + # ----------------------------------------------------------------------------- -# Test suite for BaseGenerativeFlow +# Test suite for BaseGenerativeFlow (adjusted) # ----------------------------------------------------------------------------- class TestBaseGenerativeFlow: """Tests for the abstract BaseGenerativeFlow class.""" @@ -106,21 +146,16 @@ def test_init_defaults(self, mock_dims_registry, mock_data_manager): assert flow._match_fn is None assert flow._noise_sampler is None assert flow._time_sampler is None + # generate_from_noise is now forced to True because unpaired (inherited logic) assert flow.generate_from_noise is True def test_generate_from_noise_forced_when_unpaired(self, mock_dims_registry, mock_data_manager): - """When is_paired_setting=False, generate_from_noise should be forced to True.""" - flow = ConcreteGenerativeFlow( - mock_dims_registry, - mock_data_manager, - False, - generate_from_noise=False, # user tries to set False - ) - # In BaseGenerativeFlow.__init__, if not paired, generate_from_noise becomes True + """This behaviour is already tested in TestBaseMethod, keep as integration check.""" + flow = ConcreteGenerativeFlow(mock_dims_registry, mock_data_manager, False, generate_from_noise=False) assert flow.generate_from_noise is True def test_generate_from_noise_respected_when_paired(self, mock_dims_registry, mock_data_manager): - """When paired, generate_from_noise can be set to False.""" + """When paired, the value is respected (inherited logic).""" flow = ConcreteGenerativeFlow(mock_dims_registry, mock_data_manager, True, generate_from_noise=False) assert flow.generate_from_noise is False