"""Free-function loss implementations dispatched on input type.
This module is the DOP-style core of the reconstruction objective
system. It replaces the ``BaseBinaryOP`` OOP hierarchy at
``_base_binary_op.py`` (which used ``functools.singledispatchmethod``
on instance methods) with module-level ``functools.singledispatch``
free functions. Each loss family (MDS, WMDS, GeneralError) has one
free function dispatched on the ``input`` argument's type:
- ``mds_loss(input, target, config)``
- ``wmds_loss(input, target, config)``
- ``general_error_loss(input, target, config)``
The configs are the frozen dataclasses from ``_loss_config.py``.
DOP alignment
-------------
- Principle 1 (separate code from data): the operations are free
functions; the configurations are pure-data frozen dataclasses.
No class hierarchy bundles data with behaviour.
- Principle 2 (immutable transforms): the configs are frozen=True; the
loss functions return NEW tensors/arrays (no mutation of inputs).
- Principle 3 (validate at boundary, trust internally): the configs
have already been validated at construction (``__post_init__``);
these functions trust the config and do not re-validate.
- Principle 4 (illegal states unrepresentable): invalid ``term`` or
``reduction`` cannot survive past config construction.
Backwards compatibility
------------------------
The OOP classes in ``mds.py`` and ``general_error.py`` remain as 1-line
delegation shims that call these free functions. Existing user code
that does ``mds = MultiDimensionalScaling(term='abs'); mds(x, y)``
keeps working unchanged. See the v2.21.0 release notes for the
removal schedule of the OOP shims (planned for v2.25.0, DOP Phase 8).
The ``isinstance(torch.Tensor, type)`` guard at the bottom of this
module is the same one used in ``_base_binary_op.py`` line 221.
It prevents registering the torch implementation when torch is
mocked (e.g. by Sphinx autodoc) which would otherwise raise.
"""
from __future__ import annotations
import typing as t
from functools import singledispatch
import numpy as np
try:
import torch
except ImportError:
raise ImportError(
"torch is required for this module. "
"Install with: pip install gunz-cm[3dr] (GPU) or pip install gunz-cm[3dr-cpu] (CPU)"
)
from ._loss_config import (
GeneralErrorConfig,
MDSConfig,
WMDSConfig,
)
from .consts import AVAIL_NP_F, AVAIL_TORCH_F
# ---------------------------------------------------------------------------
# Multi-Dimensional Scaling
# ---------------------------------------------------------------------------
[docs]@singledispatch
def mds_loss(
input: t.Any,
target: t.Any,
config: MDSConfig,
) -> t.Any:
"""Compute the MDS loss for the given input type.
Parameters
----------
input : numpy.ndarray or torch.Tensor
The input data (predictions).
target : numpy.ndarray or torch.Tensor
The target data (ground truth).
config : MDSConfig
Frozen configuration holding ``term``, ``reduction``, ``eps``.
Returns
-------
float or torch.Tensor
The reduced loss value. Type matches ``input``.
Raises
------
NotImplementedError
If ``input`` is neither ``np.ndarray`` nor ``torch.Tensor``
(the singledispatch default implementation).
Examples
--------
>>> cfg = MDSConfig(term='abs', reduction='mean')
>>> import numpy as np
>>> mds_loss(np.array([1.0, 2.0]), np.array([4.0, 5.0]), cfg)
"""
raise NotImplementedError(
f"mds_loss: unsupported input type {type(input).__name__!r}"
)
@mds_loss.register(np.ndarray)
def _mds_loss_numpy(
input: np.ndarray,
target: np.ndarray,
config: MDSConfig,
) -> float:
"""NumPy implementation of MDS loss."""
term_f = AVAIL_NP_F[config.term]
reduction_f = AVAIL_NP_F[config.reduction]
numerator = term_f(target - input)
denominator = term_f(target) + config.eps
return reduction_f(numerator / denominator)
# Conditionally register torch implementation (skip if torch is mocked).
if isinstance(torch.Tensor, type):
@mds_loss.register(torch.Tensor)
def _mds_loss_torch(
input: "torch.Tensor",
target: "torch.Tensor",
config: MDSConfig,
) -> "torch.Tensor":
"""PyTorch implementation of MDS loss."""
term_f = AVAIL_TORCH_F[config.term]
reduction_f = AVAIL_TORCH_F[config.reduction]
numerator = term_f(target - input)
denominator = term_f(target) + config.eps
return reduction_f(numerator / denominator)
# ---------------------------------------------------------------------------
# Weighted Multi-Dimensional Scaling
# ---------------------------------------------------------------------------
[docs]@singledispatch
def wmds_loss(
input: t.Any,
target: t.Any,
config: WMDSConfig,
) -> t.Any:
"""Compute the Weighted MDS loss for the given input type.
Parameters
----------
input : numpy.ndarray or torch.Tensor
The input data.
target : numpy.ndarray or torch.Tensor
The target data.
config : WMDSConfig
Frozen configuration holding ``term``, ``reduction``,
``weight_exp``, ``eps``.
Returns
-------
float or torch.Tensor
The reduced weighted loss value.
Raises
------
NotImplementedError
If ``input`` is neither ``np.ndarray`` nor ``torch.Tensor``.
Examples
--------
>>> cfg = WMDSConfig(term='square', reduction='mean', weight_exp=1.0)
>>> import numpy as np
>>> wmds_loss(np.array([1.0, 2.0]), np.array([4.0, 5.0]), cfg)
"""
raise NotImplementedError(
f"wmds_loss: unsupported input type {type(input).__name__!r}"
)
@wmds_loss.register(np.ndarray)
def _wmds_loss_numpy(
input: np.ndarray,
target: np.ndarray,
config: WMDSConfig,
) -> float:
"""NumPy implementation of Weighted MDS loss."""
term_f = AVAIL_NP_F[config.term]
power_f = AVAIL_NP_F['power']
sum_f = AVAIL_NP_F['sum']
weights = power_f(target, config.weight_exp)
term_loss = term_f(target - input)
weighted_loss = weights * term_loss
if config.reduction == 'mean':
total_weight = sum_f(weights)
return sum_f(weighted_loss) / (total_weight + config.eps)
# reduction == 'sum'
return sum_f(weighted_loss)
if isinstance(torch.Tensor, type):
@wmds_loss.register(torch.Tensor)
def _wmds_loss_torch(
input: "torch.Tensor",
target: "torch.Tensor",
config: WMDSConfig,
) -> "torch.Tensor":
"""PyTorch implementation of Weighted MDS loss."""
term_f = AVAIL_TORCH_F[config.term]
power_f = AVAIL_TORCH_F['power']
sum_f = AVAIL_TORCH_F['sum']
weights = power_f(target, config.weight_exp)
term_loss = term_f(target - input)
weighted_loss = weights * term_loss
if config.reduction == 'mean':
total_weight = sum_f(weights)
return sum_f(weighted_loss) / (total_weight + config.eps)
# reduction == 'sum'
return sum_f(weighted_loss)
# ---------------------------------------------------------------------------
# General Error
# ---------------------------------------------------------------------------
[docs]@singledispatch
def general_error_loss(
input: t.Any,
target: t.Any,
config: GeneralErrorConfig,
) -> t.Any:
"""Compute the general error loss for the given input type.
Parameters
----------
input : numpy.ndarray or torch.Tensor
The input data.
target : numpy.ndarray or torch.Tensor
The target data.
config : GeneralErrorConfig
Frozen configuration holding ``term`` (one of
``{l1, l2, abs, square}``; ``abs`` and ``square`` are normalised
to ``l1`` / ``l2`` at config construction) and ``reduction``.
Returns
-------
float or torch.Tensor
The reduced error value.
Examples
--------
>>> cfg = GeneralErrorConfig(term='square', reduction='mean')
>>> import numpy as np
>>> general_error_loss(np.array([1.0, 2.0]), np.array([4.0, 5.0]), cfg)
"""
raise NotImplementedError(
f"general_error_loss: unsupported input type {type(input).__name__!r}"
)
@general_error_loss.register(np.ndarray)
def _general_error_numpy(
input: np.ndarray,
target: np.ndarray,
config: GeneralErrorConfig,
) -> float:
"""NumPy implementation of general error loss."""
term_f = AVAIL_NP_F[config.term]
reduction_f = AVAIL_NP_F[config.reduction]
loss = term_f(target - input)
return reduction_f(loss)
if isinstance(torch.Tensor, type):
@general_error_loss.register(torch.Tensor)
def _general_error_torch(
input: "torch.Tensor",
target: "torch.Tensor",
config: GeneralErrorConfig,
) -> "torch.Tensor":
"""PyTorch implementation of general error loss.
PyTorch's F.l1_loss and F.mse_loss combine the term and
reduction steps; we pass the reduction string directly.
"""
term_f = AVAIL_TORCH_F[config.term]
return term_f(input, target, reduction=config.reduction)
__all__ = [
"mds_loss",
"wmds_loss",
"general_error_loss",
]