Refactor: Handling Unrecognized Arguments for Prepare

[selmanozleyen]

These are the things I want to do first. related: #596

From reading through ottjax only the __call__ functions on solvers absorb kwargs that aren't recognized. And they are because either they are given to create_geometry methods or create_initializer methods:

• https://github.com/ott-jax/ott/blob/458e265a1282d417f87bc0e07a5b6b77ffe9ae50/src/ott/initializers/quadratic/initializers.py#L134-L155
• https://github.com/ott-jax/ott/blob/458e265a1282d417f87bc0e07a5b6b77ffe9ae50/src/ott/initializers/linear/initializers_lr.py#L232-L243

So this is kind of good news because from what I read from moscot, call_fn arguments are filtered before here:

moscot/src/moscot/base/solver.py

Lines 145 to 164 in e21677e

	@classmethod
	def _partition_kwargs(cls, **kwargs: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]:
	"""Partition keyword arguments.
	Used by the :meth:`~moscot.problems.base.BaseProblem.solve`.
	Parameters
	----------
	kwargs
	Keyword arguments to partition.
	Returns
	-------
	Keyword arguments for :class:`BaseSolver` and :meth:`__call__`, respectively.
	"""
	call_kws, shared_kws = cls._call_kwargs()
	init_kwargs = {k: v for k, v in kwargs.items() if k not in call_kws or k in shared_kws}
	call_kwargs = {k: v for k, v in kwargs.items() if k in call_kws or k in shared_kws}
	return init_kwargs, call_kwargs

Overall this means:

• We can fix theprepare methods absorbing kwargs ourselves in moscot, by simply removing the usage of kwargs in concrete classes, so I will dedicate this PR to it
• But for solve we will need to refactor things from ottjax. It won't be too hard and I already spotted where they are being ignored. If we refactor ottjax we might not need any changes in moscot since the error will be thrown from there.

Related issues (will be closed after merge):

• _spatial_norm_callback implementation #698
• test_fgw_problem.py: test_prepare_costs alpha given to prepare #699
• Buggy Test: test_gw_problem gives joint_attr #700
• Buggy Test: scale_cost given to prepare in tests/problems/time/test_mixins.py #701

I separated some issues from this PR so that we can work on them later:

• Refactor: Maybe GW should inherit from FGW instead of vice versa #716 : I did this one because when inheriting docstring GW removed the arguments for linear, and this lead to duplicate docstrings
• Refactor: Check if both x and x_callback is given #715 : Because although in terms of software, current PR makes sense, there might be some cases where a parent classes argument is ignored/fixed but is not yet discovered. I am not confident about this as I am not completely familiar with the Biology use cases of all classes.
• check for warnings thrown by us #721

Now that the tests pass these are the things to do:

• clean up code (most likely requires refactoring of handle_joint_attr)
• update docstrings
• clean the TODO's and ensure that everyone agrees on the new docs and docstrings
• write tests to ensure it fails (not sure if tests are necessary at this point since the checks are static now, will probably figure out cases to test when working on Refactor: Check if both x and x_callback is given #715)

[giovp]

fantastic work @selmanozleyen !

[MUCDK]

Really great work Selman! Thanks so much!

[selmanozleyen]

@MUCDK the problem is even though 3.10 passes 3.9 now fails. The results change between these versions

[MUCDK]

Ok can we skip the test for python 3.9?

[selmanozleyen]

Sure will update the tests in a new PR then