80 lines
4.0 KiB
Python
80 lines
4.0 KiB
Python
"""Внешние контрагенты в ОСВ (спека §5.3): сколько и кто, для каждой пары (компания, счёт).
|
||
|
||
Число `n` — отрицательное биномиальное (гамма-пуассоновская смесь с дисперсией ≈ 2 × матожидание), матожидание
|
||
`mean × w(счёт) / среднее_w`, так что среднее `n` по файлам счетов 58/60/62/66/67/76 ≈ `osv_counterparties_mean`.
|
||
Потолки: 58/60/62/76 — словарь без банков и компаний группы, 66/67 — банки. Недобор из-за потолка перераспределяется на
|
||
остальные файлы того же счёта (решение 9); если потолок всё же срезал выборку — `capped=True`.
|
||
"""
|
||
|
||
import math
|
||
import random
|
||
from dataclasses import dataclass, field
|
||
|
||
from synthgen.config import Config
|
||
from synthgen.dictionary import Org
|
||
from synthgen.group import Group
|
||
from synthgen.rng import rng_for
|
||
|
||
OSV_ACCOUNTS = ("01", "58", "60", "62", "66", "67", "76", "80", "90")
|
||
WEIGHTS = {"60": 1.6, "62": 1.6, "76": 1.0, "66": 0.1, "67": 0.05, "58": 0.1}
|
||
BANK_ACCOUNTS = ("66", "67")
|
||
SERVICE_ACCOUNTS = ("60", "62", "76") # при mean = 0 в этих файлах один служебный внешний контрагент
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CounterpartyPlan:
|
||
by_file: dict[tuple[str, str], list[Org]] = field(default_factory=dict) # (код компании, корень счёта) → контрагенты
|
||
capped: bool = False # потолок словаря срезал выборку хотя бы в одном файле
|
||
|
||
|
||
def _poisson(rng: random.Random, lam: float) -> int:
|
||
if lam <= 0:
|
||
return 0
|
||
if lam > 30: # нормальное приближение: точное для нашей задачи, а цикл Кнута на больших λ медленный
|
||
return max(0, round(rng.gauss(lam, math.sqrt(lam))))
|
||
limit, k, p = math.exp(-lam), 0, 1.0
|
||
while True:
|
||
p *= rng.random()
|
||
if p <= limit:
|
||
return k
|
||
k += 1
|
||
|
||
|
||
def _draw_n(rng: random.Random, mu: float) -> int:
|
||
"""Отрицательное биномиальное с матожиданием `mu` и дисперсией 2·mu: Poisson(Gamma(mu, 1))."""
|
||
if mu <= 0:
|
||
return 0
|
||
return _poisson(rng, rng.gammavariate(mu, 1.0))
|
||
|
||
|
||
def plan(cfg: Config, group: Group) -> CounterpartyPlan:
|
||
codes = [m.code for m in group.members]
|
||
avg_w = sum(WEIGHTS.values()) / len(WEIGHTS)
|
||
capped = False
|
||
by_file: dict[tuple[str, str], list[Org]] = {}
|
||
for acct, w in WEIGHTS.items():
|
||
pool = group.banks if acct in BANK_ACCOUNTS else group.externals
|
||
cap = len(pool)
|
||
mu = cfg.osv_counterparties_mean * w / avg_w
|
||
want = [_draw_n(rng_for(cfg.seed, "cp-count", code, acct), mu) for code in codes]
|
||
if cfg.osv_counterparties_mean == 0 and acct in SERVICE_ACCOUNTS:
|
||
want = [1] * len(codes) # решение 41: при mean = 0 расчёты идут через одного служебного внешнего контрагента
|
||
got = [min(n, cap) for n in want]
|
||
shortfall = sum(want) - sum(got)
|
||
if shortfall > 0:
|
||
capped = True
|
||
room = [i for i, n in enumerate(got) if n < cap]
|
||
while shortfall > 0 and room: # круговое доведение недобора до потолка файлов, где есть место
|
||
nxt = []
|
||
for i in room:
|
||
if shortfall == 0:
|
||
break
|
||
got[i] += 1
|
||
shortfall -= 1
|
||
if got[i] < cap:
|
||
nxt.append(i)
|
||
room = nxt
|
||
for code, n in zip(codes, got):
|
||
by_file[(code, acct)] = rng_for(cfg.seed, "cp-set", code, acct).sample(pool, n) if n else []
|
||
return CounterpartyPlan(by_file=by_file, capped=capped)
|