Add predict kwargs in validation step

src/cellflow/model/_cellflow.py

@@ -193,6 +193,7 @@ def prepare_validation_data(
         name: str,
         n_conditions_on_log_iteration: int | None = None,
         n_conditions_on_train_end: int | None = None,
+        predict_kwargs: dict[str, Any] | None = None,
     ) -> None:
         """Prepare the validation data.
 
@@ -209,6 +210,10 @@ def prepare_validation_data(
         n_conditions_on_train_end
             Number of conditions to use for computation callbacks at the end of training.
             If :obj:`None`, use all conditions.
+        predict_kwargs
+            Keyword arguments for the prediction function
+            :func:`cellflow.solvers._otfm.OTFlowMatching.predict` or
+            :func:`cellflow.solvers._genot.GENOT.predict` used during validation.
 
         Returns
         -------
@@ -227,6 +232,7 @@ def prepare_validation_data(
             n_conditions_on_train_end=n_conditions_on_train_end,
         )
         self._validation_data[name] = val_data
+        self._validation_data["predict_kwargs"] = predict_kwargs
 
     def prepare_model(
         self,
@@ -488,7 +494,10 @@ def prepare_model(
             )
         else:
             raise NotImplementedError(f"Solver must be an instance of OTFlowMatching or GENOT, got {type(self.solver)}")
-        self._trainer = CellFlowTrainer(solver=self.solver)  # type: ignore[arg-type]
+        if "predict_kwargs" in self.validation_data:
+            self._trainer = CellFlowTrainer(solver=self.solver, predict_kwargs=self.validation_data["predict_kwargs"])  # type: ignore[arg-type]
+        else:
+            self._trainer = CellFlowTrainer(solver=self.solver)  # type: ignore[arg-type]
 
     def train(
         self,
@@ -543,7 +552,7 @@ def train(
             self._dataloader = OOCTrainSampler(data=self.train_data, batch_size=batch_size)
         else:
             self._dataloader = TrainSampler(data=self.train_data, batch_size=batch_size)
-        validation_loaders = {k: ValidationSampler(v) for k, v in self.validation_data.items()}
+        validation_loaders = {k: ValidationSampler(v) for k, v in self.validation_data.items() if k != "predict_kwargs"}
 
         self._solver = self.trainer.train(
             dataloader=self._dataloader,

src/cellflow/training/_trainer.py

@@ -19,8 +19,12 @@ class CellFlowTrainer:
         dataloader
             Data sampler.
         solver
-            :class:`~cellflow.solvers.OTFlowMatching` solver or :class:`~cellflow.solvers.GENOT`
-            solver with a conditional velocity field.
+            :class:`~cellflow.solvers._otfm.OTFlowMatching` or
+            :class:`~cellflow.solvers._genot.GENOT` solver with a conditional velocity field.
+        predict_kwargs
+            Keyword arguments for the prediction functions
+            :func:`cellflow.solvers._otfm.OTFlowMatching.predict` or
+            :func:`cellflow.solvers._genot.GENOT.predict` used during validation.
         seed
             Random seed for subsampling validation data.
 
@@ -32,12 +36,14 @@ class CellFlowTrainer:
     def __init__(
         self,
         solver: _otfm.OTFlowMatching | _genot.GENOT,
+        predict_kwargs: dict[str, Any] | None = None,
         seed: int = 0,
     ):
         if not isinstance(solver, (_otfm.OTFlowMatching | _genot.GENOT)):
             raise NotImplementedError(f"Solver must be an instance of OTFlowMatching or GENOT, got {type(solver)}")
 
         self.solver = solver
+        self.predict_kwargs = predict_kwargs or {}
         self.rng_subsampling = np.random.default_rng(seed)
         self.training_logs: dict[str, Any] = {}
 
@@ -61,7 +67,9 @@ def _validation_step(
             condition = batch.get("condition", None)
             true_tgt = batch["target"]
             valid_source_data[val_key] = src
-            valid_pred_data[val_key] = self.solver.predict(src, condition=condition, batched=True)
+            valid_pred_data[val_key] = self.solver.predict(
+                src, condition=condition, batched=True, **self.predict_kwargs
+            )
             valid_true_data[val_key] = true_tgt
 
         return valid_source_data, valid_true_data, valid_pred_data

tests/trainer/test_trainer.py

@@ -1,3 +1,6 @@
+import time
+
+import diffrax
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -155,3 +158,70 @@ def test_cellflow_trainer_with_custom_callback(self, dataloader, valid_loader):
         assert isinstance(logs["diff"][0], np.ndarray)
         assert logs["diff"][0].shape == (5,)
         assert 0 < np.mean(logs["diff"][0]) < 10
+
+    def test_predict_kwargs_iter(self, dataloader, valid_loader):
+        opt_1 = optax.adam(1e-3)
+        opt_2 = optax.adam(1e-3)
+        vf_1 = cellflow.networks.ConditionalVelocityField(
+            output_dim=5,
+            max_combination_length=2,
+            condition_embedding_dim=12,
+            hidden_dims=(32, 32),
+            decoder_dims=(32, 32),
+        )
+        vf_2 = cellflow.networks.ConditionalVelocityField(
+            output_dim=5,
+            max_combination_length=2,
+            condition_embedding_dim=12,
+            hidden_dims=(32, 32),
+            decoder_dims=(32, 32),
+        )
+        model_1 = _otfm.OTFlowMatching(
+            vf=vf_1,
+            match_fn=match_linear,
+            probability_path=dynamics.ConstantNoiseFlow(0.0),
+            optimizer=opt_1,
+            conditions=cond,
+            rng=vf_rng,
+        )
+        model_2 = _otfm.OTFlowMatching(
+            vf=vf_2,
+            match_fn=match_linear,
+            probability_path=dynamics.ConstantNoiseFlow(0.0),
+            optimizer=opt_2,
+            conditions=cond,
+            rng=vf_rng,
+        )
+
+        metric_to_compute = "e_distance"
+        metrics_callback = cellflow.training.Metrics(metrics=[metric_to_compute])
+
+        predict_kwargs_1 = {"stepsize_controller": diffrax.PIDController(rtol=0.1, atol=0.1)}
+        predict_kwargs_2 = {"stepsize_controller": diffrax.PIDController(rtol=1e-5, atol=1e-5)}
+
+        trainer_1 = cellflow.training.CellFlowTrainer(solver=model_1, predict_kwargs=predict_kwargs_1)
+        trainer_2 = cellflow.training.CellFlowTrainer(solver=model_2, predict_kwargs=predict_kwargs_2)
+
+        start_1 = time.time()
+        trainer_1.train(
+            dataloader=dataloader,
+            valid_loaders=valid_loader,
+            num_iterations=2,
+            valid_freq=1,
+            callbacks=[metrics_callback],
+        )
+        end_1 = time.time()
+        diff_1 = end_1 - start_1
+
+        start_2 = time.time()
+        trainer_2.train(
+            dataloader=dataloader,
+            valid_loaders=valid_loader,
+            num_iterations=2,
+            valid_freq=1,
+            callbacks=[metrics_callback],
+        )
+        end_2 = time.time()
+        diff_2 = end_2 - start_2
+
+        assert diff_2 - diff_1 > 5