Temporal Operators¶

@staticmethod
@jit
def _next_prob_seq_core_dense(m: jnp.ndarray, sub_seq: jnp.ndarray) -> jnp.ndarray:
    r"""JAX core for :class:`Next` on a dense kernel.

    Args:
        m: Dense Markov-chain transition matrix of shape ``(S, S)``, where
            column ``s`` is the distribution over next states from state ``s``.
        sub_seq: Subformula sequence of shape ``(T, S)``.

    Returns:
        Sequence of shape ``(T + 1, S)``. Row 0 is all zeros (the shifted
        operator is defined to have probability 0 at time 0), and rows 1..T
        contain one-step expectations of ``sub_seq``.
    """
    tail = (m.T @ sub_seq.T)
    tail = jnp.swapaxes(tail, 0, 1)
    zeros = jnp.zeros_like(tail[:1, :])
    return jnp.concatenate([zeros, tail], axis=0)

@staticmethod
@partial(jit, static_argnames=("K",))
def _next_prob_seq_core_compact(
    succ: jnp.ndarray, p: jnp.ndarray, sub_seq: jnp.ndarray, K: int
) -> jnp.ndarray:
    r"""JAX core for :class:`Next` on a compact kernel.

    Args:
        succ: Successor ids of shape ``(K, S)``.
        p: Successor probabilities of shape ``(K, S)`` aligned with ``succ``.
        sub_seq: Subformula sequence of shape ``(T, S)``.
        K: Max successors per state (static argument for JIT).

    Returns:
        Sequence of shape ``(T + 1, S)`` with row 0 all zeros.
    """
    def one_step(v):
        v_succ = v[succ]
        return jnp.sum(p * v_succ, axis=0)

    tail = jax.vmap(one_step)(sub_seq)
    zeros = jnp.zeros_like(tail[:1, :])
    return jnp.concatenate([zeros, tail], axis=0)

def _prob_seq(
    self,
    kernel: Kernel,
    vec_label_fn: np.ndarray,
    atom_dict: Dict[str, int],
    max_k: int | None = None,
) -> np.ndarray:
    r"""Compute the probability sequence for :math:`X\,\Phi`.

    Args:
        kernel: Markov-chain kernel (dense or compact).
        vec_label_fn: Vectorized labeling matrix ``(n_atoms, S)``.
        atom_dict: Mapping from atom string to row index.
        max_k: Local horizon (inclusive). If ``None``, defaults to 1.

    Returns:
        Float64 array of shape ``(max_k + 1, S)``.

    Notes:
        If ``max_k == 0``, returns a single row of zeros.
    """
    if max_k is None:
        max_k = self.bound_param

    n_states = kernel_n_states(kernel)
    if max_k == 0:
        return np.zeros((1, n_states), dtype=np.float64)

    sub_seq_np = self.subformula._prob_seq(kernel, vec_label_fn, atom_dict, max_k=max_k - 1)
    sub_seq_j = jnp.asarray(sub_seq_np, dtype=jnp.float64)

    if isinstance(kernel, tuple):
        succ, p = kernel
        succ_j = jnp.asarray(succ, dtype=jnp.int64)
        p_j = jnp.asarray(p, dtype=jnp.float64)
        seq_j = self._next_prob_seq_core_compact(succ_j, p_j, sub_seq_j, K=succ.shape[0])
    else:
        m_j = jnp.asarray(kernel, dtype=jnp.float64)
        seq_j = self._next_prob_seq_core_dense(m_j, sub_seq_j)

    return np.asarray(seq_j, dtype=np.float64)

def sat(
    self,
    kernel: Kernel,
    vec_label_fn: np.ndarray,
    atom_dict: Dict[str, int],
) -> np.ndarray:
    r"""Threshold the one-step probability for :math:`\mathbb{P}_{\ge p}[X\,\Phi]`.

    Args:
        kernel: Markov-chain kernel (dense or compact).
        vec_label_fn: Vectorized labeling matrix ``(n_atoms, S)``.
        atom_dict: Mapping from atom string to row index.

    Returns:
        Float64 array of shape ``(S,)`` in ``{0.0, 1.0}``.
    """
    probs = self._prob_seq(kernel, vec_label_fn, atom_dict, max_k=self.bound_param)[
        self.bound_param
    ]
    return (probs >= self.prob).astype(np.float64)

@staticmethod
@partial(jit, static_argnames=("max_k",))
def _until_prob_seq_core_dense(
    m: jnp.ndarray, sat1: jnp.ndarray, sat2: jnp.ndarray, max_k: int
) -> jnp.ndarray:
    r"""JAX core for bounded-until probabilities using a dense kernel.

    Args:
        m: Dense Markov-chain transition matrix of shape ``(S, S)``.
        sat1: Satisfaction mask for :math:`\Phi_1`, shape ``(S,)`` in ``{0,1}``.
        sat2: Satisfaction mask for :math:`\Phi_2`, shape ``(S,)`` in ``{0,1}``.
        max_k: Local bound :math:`B`.

    Returns:
        Float64 JAX array of shape ``(max_k + 1, S)``.
    """
    cont_mask = (1.0 - sat2) * sat1
    prob0 = sat2

    def body(prob, _):
        exp = m.T @ prob
        next_prob = sat2 + cont_mask * exp
        return next_prob, next_prob

    _, tail = jax.lax.scan(body, prob0, None, length=max_k)
    return jnp.concatenate([prob0[None, :], tail], axis=0)

@staticmethod
@partial(jit, static_argnames=("max_k", "K"))
def _until_prob_seq_core_compact(
    succ: jnp.ndarray, p: jnp.ndarray, sat1: jnp.ndarray, sat2: jnp.ndarray, max_k: int, K: int
) -> jnp.ndarray:
    r"""JAX core for bounded-until probabilities using a compact kernel.

    Args:
        succ: Successor ids ``(K, S)``.
        p: Successor probabilities ``(K, S)``.
        sat1: Satisfaction mask for :math:`\Phi_1`, shape ``(S,)`` in ``{0,1}``.
        sat2: Satisfaction mask for :math:`\Phi_2`, shape ``(S,)`` in ``{0,1}``.
        max_k: Local bound :math:`B`.
        K: Max successors per state (static argument for JIT).

    Returns:
        Float64 JAX array of shape ``(max_k + 1, S)``.
    """
    cont_mask = (1.0 - sat2) * sat1
    prob0 = sat2

    def exp_step(v):
        v_succ = v[succ]
        return jnp.sum(p * v_succ, axis=0)

    def body(prob, _):
        exp = exp_step(prob)
        next_prob = sat2 + cont_mask * exp
        return next_prob, next_prob

    _, tail = jax.lax.scan(body, prob0, None, length=max_k)
    return jnp.concatenate([prob0[None, :], tail], axis=0)

def _prob_seq(
    self,
    kernel: Kernel,
    vec_label_fn: np.ndarray,
    atom_dict: Dict[str, int],
    max_k: int | None = None,
) -> np.ndarray:
    r"""Compute the bounded-until probability sequence.

    Args:
        kernel: Markov-chain kernel (dense or compact).
        vec_label_fn: Vectorized labeling matrix ``(n_atoms, S)``.
        atom_dict: Mapping from atom string to row index.
        max_k: Local horizon (inclusive). If ``None``, defaults to
            :attr:`bound_param`.

    Returns:
        Float64 array of shape ``(max_k + 1, S)``.
    """
    if max_k is None:
        max_k = self.bound_param

    sat1_np = self.subformula_1.sat(kernel, vec_label_fn, atom_dict).astype(np.float64)
    sat2_np = self.subformula_2.sat(kernel, vec_label_fn, atom_dict).astype(np.float64)

    sat1_j = jnp.asarray(sat1_np, dtype=jnp.float64)
    sat2_j = jnp.asarray(sat2_np, dtype=jnp.float64)

    if isinstance(kernel, tuple):
        succ, p = kernel
        succ_j = jnp.asarray(succ, dtype=jnp.int64)
        p_j = jnp.asarray(p, dtype=jnp.float64)
        seq_j = self._until_prob_seq_core_compact(succ_j, p_j, sat1_j, sat2_j, max_k=max_k, K=succ.shape[0])
    else:
        m_j = jnp.asarray(kernel, dtype=jnp.float64)
        seq_j = self._until_prob_seq_core_dense(m_j, sat1_j, sat2_j, max_k=max_k)

    return np.asarray(seq_j, dtype=np.float64)

def sat(
    self,
    kernel: Kernel,
    vec_label_fn: np.ndarray,
    atom_dict: Dict[str, int],
) -> np.ndarray:
    r"""Threshold :math:`\mathbb{P}_{\ge p}[\Phi_1\ U^{\le B}\ \Phi_2]`.

    Args:
        kernel: Markov-chain kernel (dense or compact).
        vec_label_fn: Vectorized labeling matrix ``(n_atoms, S)``.
        atom_dict: Mapping from atom string to row index.

    Returns:
        Float64 array of shape ``(S,)`` in ``{0.0, 1.0}``.
    """
    probs = self._prob_seq(kernel, vec_label_fn, atom_dict, max_k=self.bound_param)[
        self.bound_param
    ]
    return (probs >= self.prob).astype(np.float64)

def _prob_seq(self, kernel, vec_label_fn, atom_dict, max_k=None):
    """Delegate probability-sequence computation to :attr:`_inner`."""
    return self._inner._prob_seq(kernel, vec_label_fn, atom_dict, max_k)

def sat(self, kernel, vec_label_fn, atom_dict):
    """Delegate probability-sequence computation to :attr:`_inner`."""
    return self._inner.sat(kernel, vec_label_fn, atom_dict)

def _prob_seq(self, kernel, vec_label_fn, atom_dict, max_k=None):
    """Delegate probability-sequence computation to :attr:`_inner`."""
    return self._inner._prob_seq(kernel, vec_label_fn, atom_dict, max_k)

def sat(self, kernel, vec_label_fn, atom_dict):
    """Delegate satisfaction evaluation to :attr:`_inner`."""
    return self._inner.sat(kernel, vec_label_fn, atom_dict)