# Copyright (c) 2023 zfit
from __future__ import annotations
import numpy as np
import tensorflow as tf
from zfit import z
from .interfaces import ZfitData, ZfitDimensional, ZfitOrderableDimensional, ZfitSpace
from ..util import ztyping
from ..util.container import convert_to_container
from ..util.exception import (
AxesIncompatibleError,
CoordinatesIncompatibleError,
CoordinatesUnderdefinedError,
IntentionAmbiguousError,
ObsIncompatibleError,
OverdefinedError,
)
[docs]
class Coordinates(ZfitOrderableDimensional):
def __init__(self, obs=None, axes=None):
obs, axes, n_obs = self._check_convert_obs_axes(obs, axes)
self._obs = obs
self._axes = axes
self._n_obs = n_obs
@staticmethod
def _check_convert_obs_axes(obs, axes):
if isinstance(obs, ZfitOrderableDimensional):
if axes is not None:
raise OverdefinedError(
f"Cannot use {obs}, a"
" ZfitOrderableDimensional as obs with axes not None"
" (currently, please open an issue if desired)"
)
coord = obs
return coord.obs, coord.axes, coord.n_obs
obs = convert_to_obs_str(obs, container=tuple)
axes = convert_to_axes(axes, container=tuple)
if obs is None:
if axes is None:
raise CoordinatesUnderdefinedError("Neither obs nor axes specified")
else:
if any(not isinstance(ax, int) for ax in axes):
raise ValueError(f"Axes have to be int. Currently: {axes}")
n_obs = len(axes)
else:
if any(not isinstance(ob, str) for ob in obs):
raise ValueError(f"Observables have to be strings. Currently: {obs}")
n_obs = len(obs)
if axes is not None:
if not len(obs) == len(axes):
raise CoordinatesIncompatibleError(
"obs and axes do not have the same length."
)
if not (obs or axes):
raise CoordinatesUnderdefinedError(
f"Neither obs {obs} nor axes {axes} are defined."
)
return obs, axes, n_obs
@property
def obs(self) -> ztyping.ObsTypeReturn:
"""Return the observables, string identifier for the coordinate system."""
return self._obs
@property
def axes(self) -> ztyping.AxesTypeReturn:
"""Return the axes, integer based identifier(indices) for the coordinate system."""
return self._axes
@property
def n_obs(self) -> int:
"""Return the number of observables, the dimensionality.
Corresponds to the last dimension.
"""
return self._n_obs
[docs]
def with_obs(
self,
obs: ztyping.ObsTypeInput | None,
allow_superset: bool = True,
allow_subset: bool = True,
) -> Coordinates:
"""Create a new instance that has ``obs``; sorted by or set or dropped.
The behavior is as follows:
* obs are already set:
* input obs are None: the observables will be dropped. If no axes are set, an error
will be raised, as no coordinates will be assigned to this instance anymore.
* input obs are not None: the instance will be sorted by the incoming obs. If axes or other
objects have an associated order (e.g. data, limits,...), they will be reordered as well.
If a strict subset is given (and allow_subset is True), only a subset will be returned.
This can be used to take a subspace of limits, data etc.
If a strict superset is given (and allow_superset is True), the obs will be sorted accordingly as
if the obs not contained in the instances obs were not in the input obs.
* obs are not set:
* if the input obs are None, the same object is returned.
* if the input obs are not None, they will be set as-is and now correspond to the already
existing axes in the object.
Args:
obs: Observables to sort/associate this instance with
allow_superset: if False and a strict superset of the own observables is given, an error
is raised.
allow_subset:if False and a strict subset of the own observables is given, an error
is raised.
Returns:
A copy of the object with the new ordering/observables
Raises:
CoordinatesUnderdefinedError: if obs is None and the instance does not have axes
ObsIncompatibleError: if ``obs`` is a superset and allow_superset is False or a subset and
allow_allow_subset is False
"""
obs = convert_to_obs_str(obs)
if obs is None: # drop obs, check if there are axes
if self.axes is None:
raise AxesIncompatibleError(
"cannot remove obs (using None) for a Space without axes"
)
new_coords = type(self)(obs=obs, axes=self.axes)
else:
obs = _convert_obs_to_str(obs)
if self.obs is None:
new_coords = type(self)(obs=obs, axes=self.axes)
else:
if not set(obs).intersection(self.obs):
raise ObsIncompatibleError(
f"The requested obs {obs} are not compatible with the current obs "
f"{self.obs}"
)
if not frozenset(obs) == frozenset(self.obs):
if not allow_superset and frozenset(obs) - frozenset(self.obs):
raise ObsIncompatibleError(
f"Obs {obs} are a superset of {self.obs}, not allowed according to flag."
)
if not allow_subset and set(self.obs) - set(obs):
raise ObsIncompatibleError(
f"Obs {obs} are a subset of {self.obs}, not allowed according to flag."
)
new_indices = self.get_reorder_indices(obs=obs)
new_obs = self._reorder_obs(indices=new_indices)
new_axes = self._reorder_axes(indices=new_indices)
new_coords = type(self)(obs=new_obs, axes=new_axes)
return new_coords
[docs]
def with_axes(
self,
axes: ztyping.AxesTypeInput | None,
allow_superset: bool = True,
allow_subset: bool = True,
) -> Coordinates:
"""Create a new instance that has ``axes``; sorted by or set or dropped.
The behavior is as follows:
* axes are already set:
* input axes are None: the axes will be dropped. If no observables are set, an error
will be raised, as no coordinates will be assigned to this instance anymore.
* input axes are not None: the instance will be sorted by the incoming axes. If obs or other
objects have an associated order (e.g. data, limits,...), they will be reordered as well.
If a strict subset is given (and allow_subset is True), only a subset will be returned. This can
be used to retrieve a subspace of limits, data etc.
If a strict superset is given (and allow_superset is True), the axes will be sorted accordingly as
if the axes not contained in the instances axes were not present in the input axes.
* axes are not set:
* if the input axes are None, the same object is returned.
* if the input axes are not None, they will be set as-is and now correspond to the already
existing obs in the object.
Args:
axes: Axes to sort/associate this instance with
allow_superset: if False and a strict superset of the own axeservables is given, an error
is raised.
allow_subset:if False and a strict subset of the own axeservables is given, an error
is raised.
Returns:
A copy of the object with the new ordering/axes
Raises:
CoordinatesUnderdefinedError: if obs is None and the instance does not have axes
AxesIncompatibleError: if ``axes`` is a superset and allow_superset is False or a subset and
allow_allow_subset is False
"""
axes = convert_to_axes(axes)
if axes is None: # drop axes
if self.obs is None:
raise CoordinatesUnderdefinedError(
"Cannot remove axes (using None) for a Space without obs"
)
new_coords = type(self)(obs=self.obs, axes=axes)
else:
axes = _convert_axes_to_int(axes)
if not self.axes and not len(axes) == len(self.obs):
raise AxesIncompatibleError(
f"Trying to set axes {axes} to object with obs {self.obs}"
)
if self.axes is None:
new_coords = type(self)(obs=self.obs, axes=axes)
else:
if not set(axes).intersection(self.axes):
raise AxesIncompatibleError(
f"The requested axes {axes} are not compatible with the current axes "
f"{self.axes}"
)
if not frozenset(axes) == frozenset(self.axes):
if not allow_superset and set(axes) - set(self.axes):
raise AxesIncompatibleError(
f"Axes {axes} are a superset of {self.axes}, not allowed according to flag."
)
if not allow_subset and set(self.axes) - set(axes):
raise AxesIncompatibleError(
f"Axes {axes} are a subset of {self.axes}, not allowed according to flag."
)
new_indices = self.get_reorder_indices(axes=axes)
new_obs = self._reorder_obs(indices=new_indices)
new_axes = self._reorder_axes(indices=new_indices)
new_coords = type(self)(obs=new_obs, axes=new_axes)
return new_coords
[docs]
def with_autofill_axes(self, overwrite: bool = False) -> ZfitOrderableDimensional:
"""Overwrite the axes of the current object with axes corresponding to range(len(n_obs)).
This effectively fills with (0, 1, 2,...) and can be used mostly when an object enters a PDF or
similar. ``overwrite`` allows to remove the axis first in case there are already some set.
.. code-block::
object.obs -> ('x', 'z', 'y')
object.axes -> None
object.with_autofill_axes()
object.obs -> ('x', 'z', 'y')
object.axes -> (0, 1, 2)
Args:
overwrite: If axes are already set, replace the axes with the autofilled ones.
If axes is already set and ``overwrite`` is False, raise an error.
Returns:
The object with the new axes
Raises:
AxesIncompatibleError: if the axes are already set and ``overwrite`` is False.
"""
if self.axes and not overwrite:
raise AxesIncompatibleError(
"overwrite is not allowed but axes are already set."
)
new_coords = type(self)(obs=self.obs, axes=range(self.n_obs))
return new_coords
def _reorder_obs(self, indices: tuple[int]) -> ztyping.ObsTypeReturn:
obs = self.obs
if obs is not None:
obs = tuple(obs[i] for i in indices)
return obs
def _reorder_axes(self, indices: tuple[int]) -> ztyping.AxesTypeReturn:
axes = self.axes
if axes is not None:
axes = tuple(axes[i] for i in indices)
return axes
[docs]
def get_reorder_indices(
self, obs: ztyping.ObsTypeInput = None, axes: ztyping.AxesTypeInput = None
) -> tuple[int]:
"""Indices that would order the instances obs as ``obs`` respectively the instances axes as ``axes``.
Args:
obs: Observables that the instances obs should be ordered to. Does not reorder, but just
return the indices that could be used to reorder.
axes: Axes that the instances obs should be ordered to. Does not reorder, but just
return the indices that could be used to reorder.
Returns:
New indices that would reorder the instances obs to be obs respectively axes.
Raises:
CoordinatesUnderdefinedError: If neither ``obs`` nor ``axes`` is given
"""
obs_none = obs is None
axes_none = axes is None
obs_is_defined = self.obs is not None and not obs_none
axes_is_defined = self.axes is not None and not axes_none
if not (obs_is_defined or axes_is_defined):
raise CoordinatesUnderdefinedError(
"Neither the `obs` (argument and on instance) nor `axes` (argument and on instance) are defined."
)
if obs_is_defined:
old, new = self.obs, [o for o in obs if o in self.obs]
else:
old, new = self.axes, [a for a in axes if a in self.axes]
new_indices = _reorder_indices(old=old, new=new)
return new_indices
[docs]
def reorder_x(
self,
x: tf.Tensor | np.ndarray,
*,
x_obs: ztyping.ObsTypeInput = None,
x_axes: ztyping.AxesTypeInput = None,
func_obs: ztyping.ObsTypeInput = None,
func_axes: ztyping.AxesTypeInput = None,
) -> ztyping.XTypeReturnNoData:
"""Reorder x in the last dimension either according to its own obs or assuming a function ordered with func_obs.
There are two obs or axes around: the one associated with this Coordinate object and the one associated with x.
If x_obs or x_axes is given, then this is assumed to be the obs resp. the axes of x and x will be reordered
according to ``self.obs`` resp. ``self.axes``.
If func_obs resp. func_axes is given, then x is assumed to have ``self.obs`` resp. ``self.axes`` and will be
reordered to align with a function ordered with ``func_obs`` resp. ``func_axes``.
Switching ``func_obs`` for ``x_obs`` resp. ``func_axes`` for ``x_axes`` inverts the reordering of x.
Args:
x: Tensor to be reordered, last dimension should be n_obs resp. n_axes
x_obs: Observables associated with x. If both, x_obs and x_axes are given, this has precedency over the
latter.
x_axes: Axes associated with x.
func_obs: Observables associated with a function that x will be given to. Reorders x accordingly and assumes
self.obs to be the obs of x. If both, ``func_obs`` and ``func_axes`` are given, this has precedency over the
latter.
func_axes: Axe associated with a function that x will be given to. Reorders x accordingly and assumes
self.axes to be the axes of x.
Returns:
The reordered array-like object
"""
x_reorder = x_obs is not None or x_axes is not None
func_reorder = func_obs is not None or func_axes is not None
if not (x_reorder ^ func_reorder):
raise ValueError(
"Either specify `x_obs/axes` or `func_obs/axes`, not both."
)
obs_defined = x_obs is not None or func_obs is not None
axes_defined = x_axes is not None or func_axes is not None
if obs_defined and self.obs:
if x_reorder:
coord_old = x_obs
coord_new = self.obs
elif func_reorder:
coord_new = func_obs
coord_old = self.obs
else:
assert False, "bug, should never be reached"
elif axes_defined and self.axes:
if x_reorder:
coord_old = x_axes
coord_new = self.axes
elif func_reorder:
coord_new = func_axes
coord_old = self.axes
else:
assert False, "bug, should never be reached"
else:
raise ValueError(
"Obs and self.obs or axes and self. axes not properly defined. Can only reorder on defined"
" coordinates."
)
if isinstance(x, ZfitData) and not (
coord_old == x.obs if obs_defined else x.axes
):
raise IntentionAmbiguousError(
"`reorder_x` is supposed to assume that the obs/axes of the given `x` are"
" either the one from the Space itself or the ones given. x is a ZfitData"
f" object that has obs/axes themselves {x.obs if obs_defined else x.axes}"
f" which do not coincied with the assumed obs/axes {coord_old}. Use"
f" x.value() to get the pure tensor out or rather sort Data accordingly"
f" (sort_by...)."
)
new_indices = _reorder_indices(old=coord_old, new=coord_new)
x = z.unstable.gather(x, indices=new_indices, axis=-1)
return x
def __eq__(self, other):
if not isinstance(other, Coordinates):
return NotImplemented
obs_equal = False
axes_equal = False
if self.obs is not None and other.obs is not None:
obs_equal = frozenset(self.obs) == frozenset(other.obs)
if self.axes is not None and other.axes is not None:
axes_equal = frozenset(self.axes) == frozenset(other.axes)
equal = obs_equal or axes_equal
return equal
def __hash__(self):
return 42 # always check with equal... maybe change in future, use dict that checks for different things.
def __repr__(self):
return f"<zfit Coordinates obs={self.obs}, axes={self.axes}"
def _convert_axes_to_int(axes):
if isinstance(axes, ZfitSpace):
axes = axes.axes
else:
axes = convert_to_container(axes, container=tuple)
return axes
def _convert_obs_to_str(obs):
if isinstance(obs, ZfitSpace):
obs = obs.obs
else:
obs = convert_to_container(obs, container=tuple)
return obs
def _reorder_indices(old: list | tuple, new: list | tuple) -> tuple[int]:
new_indices = tuple(old.index(o) for o in new)
return new_indices
def convert_to_axes(axes, container=tuple):
"""Convert `obs` to the list of obs, also if it is a :py:class:`~ZfitSpace`. Return None if axes is None.
Raises
TypeError: if the axes are not int
"""
if axes is None:
return axes
axes = convert_to_container(value=axes, container=container)
if isinstance(axes, ZfitDimensional):
new_axes = axes.axes
else:
new_axes = []
for axis in axes:
if not isinstance(axis, int):
raise TypeError(f"Axes have to be int, not {axis} as in {axes}")
else:
new_axes.append(axis)
return container(new_axes)
def convert_to_obs_str(obs, container=tuple):
"""Convert `obs` to the list of obs, also if it is a :py:class:`~ZfitSpace`. Return None if obs is None.
Raises:
TypeError: if the observable is not a string
"""
if obs is None:
return obs
if isinstance(obs, ZfitDimensional):
new_obs = obs.obs
else:
obs = convert_to_container(value=obs, container=container)
new_obs = []
for ob in obs:
if not isinstance(ob, str):
raise TypeError(f"Observables have to be string, not {ob} as in {obs}")
else:
new_obs.append(ob)
return container(new_obs)
