delay_space.py

from dataclasses import dataclass

import numpy as np

# Reduce display precision.
np.set_printoptions(precision=3)

# Ignore divide-by-zero and invalid-value errors.
np.seterr(divide="ignore", invalid="ignore")

a = np.array([0, 0.5, 1, 2, 4])
b = np.array([-5, -3, -1, -0.2, 0])

print(f"a={a}")
print(f"b={b}")

# All "Sections", "Equations", and "Tables" in this file are references to
# `Temporal_Logic_Addendum.pdf`.
print("\nSection 2: Convert non-negative numbers between importance- and delay-space")
print("-------------------------------------------------------------------------------")


def to_delay_space_partial(x: np.ndarray) -> np.ndarray:
    """Equation 2: Convert ``x`` from importance-space to delay-space. ``x`` must be
    non-negative.
    """
    assert np.all(x >= 0)
    return np.where(x > 0, -np.log(x), np.inf)


# Convention: Delay-space variables have a trailing underscore.
a_ = to_delay_space_partial(a)

print("\nConverting a to delay-space:")
print(f"a_={a_}")


def to_importance_space_partial(x_: np.ndarray) -> np.ndarray:
    """Equation 1: Convert ``x_`` from delay-space to importance-space."""
    return np.exp(-x_)


roughly_a = to_importance_space_partial(a_)
np.testing.assert_allclose(roughly_a, a)

print("\nConverting a_ back to importance-space:")
print(f"roughly_a={roughly_a}")

print("\nSection 3: Convert negative numbers between importance- and delay-space")
print("-------------------------------------------------------------------------------")


@dataclass
class PosNeg:
    """Section 3: PosNeg represents an ``ndarray`` as separate ``pos`` and ``neg``
    components.

    ``ndarray == pos - neg``

    In importance-space, at least one of ``pos`` or ``neg`` must be zero.

    In delay-space, at least one of ``pos`` or ``neg`` must be infinity.
    """

    pos: np.ndarray
    neg: np.ndarray

    def __repr__(self) -> str:
        return f"<pos={self.pos}, neg={self.neg}>"


def to_delay_space_full(x: np.ndarray) -> PosNeg:
    """Equations 5, 6: Convert ``x`` from importance-space to delay-space. ``x`` may be
    negative.
    """
    return PosNeg(
        pos=to_delay_space_partial(np.maximum(x, 0)),
        neg=to_delay_space_partial(np.abs(np.minimum(x, 0))),
    )


b_ = to_delay_space_full(b)

print("\nConverting b to delay-space:")
print(f"b_={b_}")


def to_importance_space_full(x_: PosNeg) -> np.ndarray:
    """Convert ``x_`` from delay-space to importance-space. ``x_`` may be negative in
    importance-space."""
    return to_importance_space_partial(x_.pos) - to_importance_space_partial(x_.neg)


roughly_b = to_importance_space_full(b_)
np.testing.assert_allclose(roughly_b, b)

print("\nConverting b_ back to importance-space:")
print(f"roughly_b={roughly_b}")

print("\nSection 5.1: Delay-space multiplication")
print("-------------------------------------------------------------------------------")


def delay_space_multiply(x_: PosNeg, y_: PosNeg) -> PosNeg:
    """Equations 10, 11: Multiply two delay-space values. Return a delay-space
    product."""
    return PosNeg(
        pos=np.minimum(x_.pos + y_.pos, x_.neg + y_.neg),
        neg=np.minimum(x_.pos + y_.neg, x_.neg + y_.pos),
    )


a_ = to_delay_space_full(a)
product_ = delay_space_multiply(a_, b_)

print("\nMultiplying a_ and b_:")
print(f"product_={product_}")

product = to_importance_space_full(product_)
np.testing.assert_allclose(product, a * b)

print("\nConverting product_ back to importance-space:")
print(f"expected_product={a * b}")
print(f"converted_product={product}")

print("\nSection 5.2: Delay-space addition")
print("-------------------------------------------------------------------------------")


def nLSE(x_: np.ndarray, y_: np.ndarray) -> np.ndarray:
    """Equation 8: Negative log-sum-exp. A partial implementation of delay-space
    addition that requires non-negative inputs.
    """
    return -np.log(np.exp(-x_) + np.exp(-y_))


def nLDE(x_: np.ndarray, y_: np.ndarray) -> np.ndarray:
    """Equation 9: Negative log-difference-exp. A partial implementation of delay-space
    subtraction that requires non-negative inputs and outputs.
    """
    return -np.log(np.exp(-x_) - np.exp(-y_))


def delay_space_add(x_: PosNeg, y_: PosNeg, nlse=nLSE, nlde=nLDE) -> PosNeg:
    """Equation 15: Add two delay-space values. Return a delay-space sum."""

    def nonneg_sub(pos_sum_: np.ndarray, neg_sum_: np.ndarray) -> PosNeg:
        """Equations 13, 14: Another partial implementation of delay-space subtraction
        that requires non-negative inputs, but allows negative output.
        """
        return PosNeg(
            pos=np.where(pos_sum_ >= neg_sum_, np.inf, nlde(pos_sum_, neg_sum_)),
            neg=np.where(pos_sum_ > neg_sum_, nlde(neg_sum_, pos_sum_), np.inf),
        )

    return nonneg_sub(pos_sum_=nlse(x_.pos, y_.pos), neg_sum_=nlse(x_.neg, y_.neg))


sum_ = delay_space_add(a_, b_)

print("\nAdding a_ and b_:")
print(f"sum_={sum_}")

sum = to_importance_space_full(sum_)
np.testing.assert_allclose(sum, a + b)

print("\nConverting sum_ back to importance-space:")
print(f"expected_sum={a + b}")
print(f"converted_sum={sum}")

print("\nSection 5.3: Delay-space subtraction")
print("-------------------------------------------------------------------------------")


def delay_space_negate(x_: PosNeg) -> PosNeg:
    """Equations 16, 17: Negation just swaps ``pos`` and ``neg``."""
    return PosNeg(pos=x_.neg, neg=x_.pos)


def delay_space_subtract(x_: PosNeg, y_: PosNeg) -> PosNeg:
    """Equation 18: Delay-space subtraction just composes delay-space addition and
    delay-space negation.
    """
    return delay_space_add(x_, delay_space_negate(y_))


difference_ = delay_space_subtract(a_, b_)

print("\nSubtracting a_ and b_:")
print(f"difference_={difference_}")

difference = to_importance_space_full(difference_)
np.testing.assert_allclose(difference, a - b)

print("\nConverting difference_ back to importance-space:")
print(f"expected_difference={a - b}")
print(f"converted_difference={difference}")


print("\nSection 6.1: Approximating nLSE")
print("-------------------------------------------------------------------------------")


def nLSE_approx(x_: np.ndarray, y_: np.ndarray) -> np.ndarray:
    """Equation 19. C and D are approximation constants from Table 1."""
    C = np.array(
        [
            -0.8177787097533084, -1.0113960852257735, -1.239795093320612,
            -1.511997607789788, -1.8419050170277562, -2.2525262779402597,
            -2.7861121434031597, -3.534084117977515, -4.765167636010712
        ]
    )  # fmt: off
    D = np.array(
        [
            -0.6527503110331004, -0.5119890797338222, -0.39226497992800746,
            -0.29119330170774393, -0.20697346514725667, -0.13823912527430826,
            -0.08395802261005558, -0.04336412687822344, -0.015912101707151294
        ]
    )  # fmt: off

    output = []
    for u, v in zip(x_, y_):
        smaller = min(u, v)
        larger = max(u, v)

        smallest = smaller
        for c, d in zip(C, D):
            smallest = min(smallest, max(larger + c, smaller + d))

        output.append(smallest)

    return np.array(output)


nlse_actual = nLSE(a_.pos, b_.neg)
nlse_approx = nLSE_approx(a_.pos, b_.neg)

print("\nApproximating nLSE(a_.pos, b_.neg):")
print(f"nlse_actual={nlse_actual}")
print(f"nlse_approx={nlse_approx}")

print("\nSection 6.2: Approximating nLDE")
print("-------------------------------------------------------------------------------")


def inhibit(i: np.ndarray, d: np.ndarray) -> np.ndarray:
    """Equation 21: ``i`` is the inhibiting event, and ``d`` is the data event.

    If the data event occurs before the inhibiting event, pass through the data event.

    If the inhibiting event occurs before the data event, inhibit the data event. The
    data event will never occur.
    """
    return np.where(i < d, np.inf, d)


def nLDE_approx(x_: np.ndarray, y_: np.ndarray) -> np.ndarray:
    """Equation 20: E and F are approximation constants from Table 2."""
    E = np.array(
        [
            4.491404993750656, 2.5545397049202143, 1.53435312464828,
            0.7829328320184368, 0.13156826590768395, -0.504783612014241,
            -1.2014232810777856, -2.076962611740086, -3.464620771514539
        ]
    )  # fmt: off
    F = np.array(
        [
            4.491404993750656, 2.5841383795968533, 1.6618997600342356,
            1.076483980201438, 0.6716856269932849, 0.3855711971301814,
            0.1879530372001602, 0.06268055855831878, 0.0
        ]
    )  # fmt: off

    output = []
    for u, v in zip(x_, y_):
        output.append(min([inhibit(v + e, u + f) for e, f in zip(E, F)]))

    return np.array(output)


nlde_actual = nLDE(a_.pos, b_.neg)
nlde_approx = nLDE_approx(a_.pos, b_.neg)

print("\nApproximating nLDE(a_.pos, b_.neg):")
print(f"nlde_actual={nlde_actual}")
print(f"nlde_approx={nlde_approx}")

print("\nUse nLDE and nLSE approximations in delay-space addition")
print("-------------------------------------------------------------------------------")

approx_sum_ = delay_space_add(a_, b_, nlse=nLSE_approx, nlde=nLDE_approx)

print("\nAdding a_ and b_ with nLSE and nLDE approximations:")
print(f"actual_sum_={sum_}")
print(f"approx_sum_={approx_sum_}")

approx_sum = to_importance_space_full(approx_sum_)
np.testing.assert_allclose(approx_sum, a + b, atol=1)

print("\nConverting approx_sum_ back to importance-space:")
print(f"actual_sum={a + b}")
print(f"approx_sum={approx_sum}")