ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein#
- class ott.solvers.quadratic.gromov_wasserstein.GromovWasserstein(*args, warm_start=None, unscale_last_linearization=False, quad_initializer=None, progress_fn=None, kwargs_init=None, **kwargs)[source]#
Gromov-Wasserstein solver [Peyré et al., 2016].
- Parameters:
args (
Any
) – Positional arguments forWassersteinSolver
.warm_start (
Optional
[bool
]) – Whether to initialize (low-rank) Sinkhorn calls using values from the previous iteration. If None, warm starts are not used for standard Sinkhorn, but used for low-rank Sinkhorn.unscale_last_linearization (
bool
) – Whether to remove any scaling from the cost matrices of the last linearization stored ingeom
. This has the practical benefit that, while the OT coupling matrices obtained with GW might have been computed by re-scaling cost matrices for numerical stability, the last linearization stored in the geometry will be unscaled and recomputed with the original cost values.quad_initializer (
Union
[Literal
['random'
,'rank2'
,'k-means'
,'generalized-k-means'
],BaseQuadraticInitializer
,None
]) – Quadratic initializer. If the solver is entropic,QuadraticInitializer
is always used. Otherwise, the quadratic initializer wraps the low-rank Sinkhorn initializers. If None, the low-rank initializer will be selected in a problem-specific manner. If bothgeom_xx
andgeom_yy
arePointCloud
orLRCGeometry
, useKMeansInitializer
. Otherwise, useRandomInitializer
.progress_fn (
Optional
[Callable
[[Tuple
[ndarray
,ndarray
,ndarray
,GWState
]],None
]]) – callback function which gets called during the Gromov-Wasserstein iterations, so the user can display the error at each iteration, e.g., using a progress bar. Seedefault_progress_fn()
for a basic implementation.kwargs_init (
Optional
[Mapping
[str
,Any
]]) – Keyword arguments when creating the initializer.kwargs (
Any
) – Keyword arguments forWassersteinSolver
.
Methods
create_initializer
(prob)Create quadratic, possibly low-rank initializer.
init_state
(prob, init, rng)Initialize the state of the Gromov-Wasserstein iterations.
output_from_state
(state)Create an output from a loop state.
Attributes
Whether the solver is low-rank.
Whether to initialize (low-rank) Sinkhorn using previous solutions.