mythos.observables.wasserstein
==============================

.. py:module:: mythos.observables.wasserstein

.. autoapi-nested-parse::

   Wasserstein distance observables.



Classes
-------

.. autoapisummary::

   mythos.observables.wasserstein.WassersteinDistance
   mythos.observables.wasserstein.WassersteinDistanceMapped


Functions
---------

.. autoapisummary::

   mythos.observables.wasserstein.wasserstein_1d
   mythos.observables.wasserstein._compute_wasserstein_distance


Module Contents
---------------

.. py:function:: wasserstein_1d(u: mythos.utils.types.Arr_N, v: mythos.utils.types.Arr_N, u_weights: mythos.utils.types.Arr_N | None = None, v_weights: mythos.utils.types.Arr_N | None = None) -> mythos.utils.types.Scalar

   Compute the 1D Wasserstein distance between two distributions u and v.


.. py:function:: _compute_wasserstein_distance(obs_values: mythos.utils.types.Arr_N, v: mythos.utils.types.Arr_N, weights: mythos.utils.types.Arr_N | None = None, v_weights: mythos.utils.types.Arr_N | None = None) -> mythos.utils.types.Scalar

.. py:class:: WassersteinDistance

   Compute the 1D Wasserstein distance between two distributions.

   The U distribution is obtained by calling the supplied observable on the
   trajectory, and the V distribution is provided as a fixed reference
   distribution. Weights can optionally be provided for the V distribution as a
   property, and for the U distribution at call time.

   The observable, when called on a trajectory, should return a (n_states,
   n_values) array, where n_states is the number of states in the trajectory.
   This will be flattened on its way into the Wasserstein distance computation.

   The weights supplied to the call method are expected to correspond to states
   in the trajectory, and will apply to all values in the observable output
   distribution for that state.

   .. attribute:: observable

      The observable whose output distribution defines U.

   .. attribute:: v_distribution

      The fixed reference distribution V to compare against.

   .. attribute:: v_weights

      Optional weights for the V distribution (should sum to 1).


   .. py:attribute:: observable
      :type:  mythos.observables.base.BaseObservable


   .. py:attribute:: v_distribution
      :type:  mythos.utils.types.Arr_N


   .. py:attribute:: v_weights
      :type:  mythos.utils.types.Arr_N | None
      :value: None



   .. py:method:: __call__(trajectory: mythos.simulators.io.SimulatorTrajectory, weights: mythos.utils.types.Arr_N | None = None) -> mythos.utils.types.Scalar

      Compute the Wasserstein distance between observable and reference distributions.



.. py:class:: WassersteinDistanceMapped

   Compute the 1D Wasserstein distance between two distributions, by key.

   This is a generalization of WassersteinDistance that allows computing
   distances for multiple observables and reference distributions at once, by
   key. The input observable is expected to return a dictionary mapping keys to
   observable outputs, and the v_distribution_map (value corresponding to
   v_distribution of `:class:WassersteinDistanceMapped`) and v_weights_map
   (value corresponding to v_weights of `:class:WassersteinDistanceMapped`)
   should have matching keys.

   See `:class:WassersteinDistance` for more information on inputs and calling.

   .. attribute:: observable

      The observable whose output distribution defines U, expected
      to return a dictionary mapping keys to observable outputs.

   .. attribute:: v_distribution_map

      Dictionary mapping keys to fixed reference
      distributions V to compare against.

   .. attribute:: v_weights_map

      Optional dictionary mapping keys to weights for the V
      distributions.


   .. py:attribute:: observable
      :type:  mythos.observables.base.BaseObservable


   .. py:attribute:: v_distribution_map
      :type:  dict[str, mythos.utils.types.Arr_N]


   .. py:attribute:: v_weights_map
      :type:  dict[str, mythos.utils.types.Arr_N | None]


   .. py:method:: __call__(trajectory: mythos.simulators.io.SimulatorTrajectory, weights: mythos.utils.types.Arr_N | None = None) -> dict[str, mythos.utils.types.Scalar]

      Compute the Wasserstein distance between all observable and reference distributions, by key.