src/iohmm_evac/dgp/emissions.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (C) 2026 SWGY, Inc.
"""State-conditional emission model.

Three emission channels, sampled per household per step:

* ``D`` — noisy departure indicator (Bernoulli)
* ``X`` — displacement from home in km (Gaussian / half-Gaussian)
* ``C`` — communication-activity count (Poisson)
"""

from __future__ import annotations

import numpy as np
from numpy.random import Generator

from iohmm_evac.params import EmissionParams
from iohmm_evac.types import BoolArray, FloatArray, IntArray, State

__all__ = ["sample_emissions"]


def _state_lambda(emissions: EmissionParams) -> FloatArray:
    """Per-state Poisson rate vector, indexed by ``State``."""
    return np.array(
        [
            emissions.lambda_ua,
            emissions.lambda_aw,
            emissions.lambda_pr,
            emissions.lambda_er,
            emissions.lambda_sh,
        ],
        dtype=np.float64,
    )


def _state_departure_p(emissions: EmissionParams) -> FloatArray:
    """Per-state Bernoulli rate for the departure indicator."""
    p = np.full(State.n_states(), emissions.p_departure_other, dtype=np.float64)
    p[int(State.ER)] = emissions.p_departure_er
    return p


def sample_emissions(
    state: IntArray,
    evac_path: IntArray,
    tir: FloatArray,
    destination: FloatArray,
    congestion_t: float,
    emissions: EmissionParams,
    rng: Generator,
) -> tuple[BoolArray, FloatArray, IntArray]:
    """Sample (departure, displacement, comm-count) for every household.

    Vectorized over households. ``evac_path`` is encoded as 0=NONE, 1=AWAY,
    2=HOME.
    """
    n = state.shape[0]
    p_vec = _state_departure_p(emissions)
    lam_vec = _state_lambda(emissions)

    # Departure: Bernoulli with state-conditioned rate.
    p_per_hh = p_vec[state]
    departure = rng.random(size=n) < p_per_hh

    # Communication: Poisson with state-conditioned rate.
    comm = rng.poisson(lam_vec[state], size=n).astype(np.int64)

    # Displacement: branch by state and (for SH) by evac_path.
    displacement = np.empty(n, dtype=np.float64)

    idle_sigma = emissions.displacement_idle_sigma
    route_sigma = emissions.displacement_route_sigma
    dest_sigma = emissions.displacement_destination_sigma

    pre_route = (state == State.UA) | (state == State.AW) | (state == State.PR)
    if pre_route.any():
        displacement[pre_route] = np.abs(rng.normal(0.0, idle_sigma, size=int(pre_route.sum())))

    en_route = state == State.ER
    if en_route.any():
        v_eff = emissions.v_free * (1.0 - emissions.congestion_penalty * congestion_t)
        # Distance traveled is min(tir * v_eff, dest_i) plus N(0, route_sigma^2).
        progress = np.minimum(tir[en_route] * v_eff, destination[en_route])
        displacement[en_route] = progress + rng.normal(0.0, route_sigma, size=int(en_route.sum()))

    sheltered = state == State.SH
    if sheltered.any():
        away = sheltered & (evac_path == 1)
        home = sheltered & (evac_path != 1)
        if away.any():
            displacement[away] = rng.normal(destination[away], dest_sigma, size=int(away.sum()))
        if home.any():
            displacement[home] = np.abs(rng.normal(0.0, idle_sigma, size=int(home.sum())))

    return departure, displacement, comm