mergeron 2025.739439.11__py3-none-any.whl → 2025.739439.12__py3-none-any.whl

mergeron/gen/data_generation_functions.py

@@ -4,13 +4,12 @@ from __future__ import annotations
 
 from typing import Literal
 
-import numexpr as ne  # type: ignore
 import numpy as np
 from attrs import evolve
 from numpy.random import SeedSequence
 
 from .. import (  # noqa: TID252
-    DEFAULT_REC_RATIO,
+    DEFAULT_REC,
     NTHREADS,
     VERSION,
     ArrayDouble,
@@ -27,14 +26,14 @@ from ..core.pseudorandom_numbers import (  # noqa: TID252
 from . import (
     DEFAULT_BETA_BND_DIST_PARMS,
     DEFAULT_FCOUNT_WTS,
-    MarginSampleData,
+    MarginsData,
+    MarketSharesData,
+    MarketShareSpec,
     PCMDistribution,
     PCMRestriction,
     PCMSpec,
-    PriceSampleData,
+    PricesData,
     PriceSpec,
-    ShareSampleData,
-    ShareSpec,
     SHRDistribution,
     SSZConstant,
 )
@@ -42,14 +41,14 @@ from . import (
 __version__ = VERSION
 
 
-def gen_share_data(
+def market_share_sampler(
     _sample_size: int,
-    _share_spec: ShareSpec,
+    _share_spec: MarketShareSpec,
     _fcount_rng_seed_seq: SeedSequence | None,
     _mktshr_rng_seed_seq: SeedSequence,
     _nthreads: int,
     /,
-) -> ShareSampleData:
+) -> MarketSharesData:
     """Generate share data.
 
     Parameters
@@ -65,7 +64,7 @@ def gen_share_data(
 
     Returns
     -------
-        Arrays representing shares, diversion ratios, etc. structured as a :ShareSampleData:
+        Arrays representing shares, diversion ratios, etc. structured as a :MarketSharesData:
 
     """
     dist_type_mktshr, firm_count_prob_wts, dist_parms_mktshr, recapture_form = (
@@ -74,12 +73,12 @@ def gen_share_data(
     )
 
     if dist_type_mktshr == SHRDistribution.UNI:
-        mkt_share_sample = gen_market_shares_uniform(
+        mkt_share_sample = market_share_sampler_uniform(
             _sample_size, dist_parms_mktshr, _mktshr_rng_seed_seq, _nthreads
         )
 
     elif dist_type_mktshr.name.startswith("DIR_"):
-        mkt_share_sample = _gen_market_shares_dirichlet_multimarket(
+        mkt_share_sample = _market_share_sampler_dirichlet_multimarket(
             _sample_size,
             recapture_form,
             dist_type_mktshr,
@@ -97,9 +96,9 @@ def gen_share_data(
 
     # If recapture_form == "inside-out", recalculate _aggregate_purchase_prob
     frmshr_array = mkt_share_sample.mktshr_array[:, :2]
-    r_bar_ = _share_spec.recapture_ratio or DEFAULT_REC_RATIO
+    r_bar_ = _share_spec.recapture_rate or DEFAULT_REC
     if recapture_form == RECForm.INOUT:
-        mkt_share_sample = ShareSampleData(
+        mkt_share_sample = MarketSharesData(
             mkt_share_sample.mktshr_array,
             mkt_share_sample.fcounts,
             mkt_share_sample.nth_firm_share,
@@ -109,13 +108,13 @@ def gen_share_data(
     return mkt_share_sample
 
 
-def gen_market_shares_uniform(
+def market_share_sampler_uniform(
     _s_size: int,
     _dist_parms_mktshr: ArrayDouble,
     _mktshr_rng_seed_seq: SeedSequence,
     _nthreads: int,
     /,
-) -> ShareSampleData:
+) -> MarketSharesData:
     """Generate merging-firm shares from Uniform distribution on the 3-D simplex.
 
     Parameters
@@ -167,12 +166,12 @@ def gen_market_shares_uniform(
     nth_firm_share_.fill(np.nan)
     aggregate_purchase_prob_.fill(np.nan)
 
-    return ShareSampleData(
+    return MarketSharesData(
         mktshr_array_, fcounts_, nth_firm_share_, aggregate_purchase_prob_
     )
 
 
-def _gen_market_shares_dirichlet_multimarket(
+def _market_share_sampler_dirichlet_multimarket(
     _s_size: int,
     _recapture_form: RECForm,
     _dist_type_dir: SHRDistribution,
@@ -182,7 +181,7 @@ def _gen_market_shares_dirichlet_multimarket(
     _mktshr_rng_seed_seq: SeedSequence,
     _nthreads: int = NTHREADS,
     /,
-) -> ShareSampleData:
+) -> MarketSharesData:
     """Dirichlet-distributed shares with multiple firm-counts.
 
     Firm-counts may be specified as having Uniform distribution over the range
@@ -202,7 +201,7 @@ def _gen_market_shares_dirichlet_multimarket(
         Whether Dirichlet is Flat or Asymmetric
 
     _recapture_form
-        r_1 = r_2 if "proportional", otherwise MNL-consistent
+        r_1 = r_2 if RECForm.FIXED, otherwise share-proportional
 
     _fcount_rng_seed_seq
         seed firm count rng, for replicable results
@@ -247,7 +246,7 @@ def _gen_market_shares_dirichlet_multimarket(
 
     if _dist_type_dir == SHRDistribution.DIR_COND:
 
-        def _gen_dir_alphas(_fcv: int) -> ArrayDouble:
+        def _alphas_builder(_fcv: int) -> ArrayDouble:
             dat_ = [2.5] * 2
             if _fcv > len(dat_):
                 dat_ += [1.0 / (_fcv - 2)] * (_fcv - 2)
@@ -255,7 +254,7 @@ def _gen_market_shares_dirichlet_multimarket(
 
     else:
 
-        def _gen_dir_alphas(_fcv: int) -> ArrayDouble:
+        def _alphas_builder(_fcv: int) -> ArrayDouble:
             return np.array(dir_alphas_full[:_fcv], float)
 
     fcounts_ = prng(_fcount_rng_seed_seq).choice(
@@ -270,10 +269,10 @@ def _gen_market_shares_dirichlet_multimarket(
     mktshr_array_ = np.empty((_s_size, fc_max), float)
     for f_val, f_sseq in zip(fcount_keys, mktshr_seed_seq_ch, strict=True):
         fcounts_match_rows = np.where(fcounts_ == f_val)[0]
-        dir_alphas_test = _gen_dir_alphas(f_val)
+        dir_alphas_test = _alphas_builder(f_val)
 
         try:
-            mktshr_sample_f = gen_market_shares_dirichlet(
+            mktshr_sample_f = market_share_sampler_dirichlet(
                 dir_alphas_test,
                 len(fcounts_match_rows),
                 _recapture_form,
@@ -306,19 +305,19 @@ def _gen_market_shares_dirichlet_multimarket(
             )
         )
 
-    return ShareSampleData(
+    return MarketSharesData(
         mktshr_array_, fcounts_, nth_firm_share_, aggregate_purchase_prob_
     )
 
 
-def gen_market_shares_dirichlet(
+def market_share_sampler_dirichlet(
     _dir_alphas: ArrayDouble,
     _s_size: int,
     _recapture_form: RECForm,
     _mktshr_rng_seed_seq: SeedSequence,
     _nthreads: int,
     /,
-) -> ShareSampleData:
+) -> MarketSharesData:
     """Dirichlet-distributed shares with fixed firm-count.
 
     Parameters
@@ -389,7 +388,7 @@ def gen_market_shares_dirichlet(
         aggregate_purchase_prob_ = 1 - mktshr_array[:, [-1]]
         mktshr_array = mktshr_array[:, :-1] / aggregate_purchase_prob_
 
-    return ShareSampleData(
+    return MarketSharesData(
         mktshr_array,
         (mktshr_array.shape[-1] * np.ones((_s_size, 1))).astype(np.int64),
         mktshr_array[:, [-1]],
@@ -397,9 +396,9 @@ def gen_market_shares_dirichlet(
     )
 
 
-def gen_divr_array(
+def diversion_ratios_builder(
     _recapture_form: RECForm,
-    _recapture_ratio: float | None,
+    _recapture_rate: float | None,
     _frmshr_array: ArrayDouble,
     _aggregate_purchase_prob: ArrayDouble,
     /,
@@ -415,8 +414,8 @@ def gen_divr_array(
     _recapture_form
         Enum specifying Fixed (proportional), Inside-out, or Outside-in
 
-    _recapture_ratio
-        If recapture is proportional or inside-out, the recapture ratio
+    _recapture_rate
+        If recapture is proportional or inside-out, the recapture rate
         for the firm with the smaller share.
 
     _frmshr_array
@@ -437,23 +436,25 @@ def gen_divr_array(
         Merging-firm diversion ratios for mergers in the sample.
 
     """
-    divr_array: ArrayDouble
+    divratio_array: ArrayDouble
+
     if _recapture_form == RECForm.FIXED:
-        _frmshr_array_rev = _frmshr_array[:, ::-1]
-        divr_array = ne.evaluate(
-            "_recapture_ratio * _frmshr_array_rev / (1 - _frmshr_array)"
+        divratio_array = (
+            (DEFAULT_REC if _recapture_rate is None else _recapture_rate)
+            * _frmshr_array[:, ::-1]
+            / (1 - _frmshr_array)
         )
 
     else:
-        purchprob_array = ne.evaluate("_aggregate_purchase_prob * _frmshr_array")
-        purchprob_array_rev = purchprob_array[:, ::-1]  # noqa: F841
-        divr_array = ne.evaluate("purchprob_array_rev / (1 - purchprob_array)")
+        _purchase_prob = np.einsum("ij,ij->ij", _aggregate_purchase_prob, _frmshr_array)
+        divratio_array = _purchase_prob[:, ::-1] / (1 - _purchase_prob)
 
     divr_assert_test = (
         (np.round(np.einsum("ij->i", _frmshr_array), 15) == 1)
-        | (np.argmin(_frmshr_array, axis=1) == np.argmax(divr_array, axis=1))
+        | (np.argmin(_frmshr_array, axis=1) == np.argmax(divratio_array, axis=1))
     )[:, None]
     if not all(divr_assert_test):
+        print(_frmshr_array, divratio_array)
         raise ValueError(
             "{} {} {} {}".format(
                 "Data construction fails tests:",
@@ -463,10 +464,10 @@ def gen_divr_array(
             )
         )
 
-    return divr_array
+    return divratio_array
 
 
-def gen_margin_price_data(
+def prices_sampler(  # noqa: PLR0914
     _frmshr_array: ArrayDouble,
     _nth_firm_share: ArrayDouble,
     _aggregate_purchase_prob: ArrayDouble,
@@ -477,13 +478,13 @@ def gen_margin_price_data(
     _pr_rng_seed_seq: SeedSequence | None,
     _nthreads: int,
     /,
-) -> tuple[MarginSampleData, PriceSampleData]:
+) -> tuple[MarginsData, PricesData]:
     """Generate margin and price data for mergers in the sample.
 
     Parameters
     ----------
     _frmshr_array
-        Merging-firm shares; see :class:`mergeron.gen.ShareSpec`.
+        Merging-firm shares; see :class:`mergeron.gen.MarketShareSpec`.
 
     _nth_firm_share
         Share of the nth firm in the sample.
@@ -517,7 +518,7 @@ def gen_margin_price_data(
     -------
         Simulated margin- and price-data arrays for mergers in the sample.
     """
-    margin_data = MarginSampleData(
+    margin_data = MarginsData(
         np.empty_like(_frmshr_array), np.ones(len(_frmshr_array)) == 0
     )
 
@@ -551,7 +552,7 @@ def gen_margin_price_data(
             # generate price and margin data
             frmshr_array_plus = np.hstack((_frmshr_array, _nth_firm_share))
             pcm_spec_here = evolve(_pcm_spec, pcm_restriction=PCMRestriction.IID)
-            margin_data = _gen_margin_data(
+            margin_data = _margins_sampler(
                 frmshr_array_plus,
                 np.ones_like(frmshr_array_plus, np.float64),
                 _aggregate_purchase_prob,
@@ -560,37 +561,41 @@ def gen_margin_price_data(
                 _nthreads,
             )
 
-            pcm_array, mnl_test_array = (
-                getattr(margin_data, _f) for _f in ("pcm_array", "mnl_test_array")
+            pcm_array, mnl_test = (
+                getattr(margin_data, _f) for _f in ("pcm_array", "mnl_test")
             )
 
-            price_array_here = ne.evaluate("1 / (1 - pcm_array)")
+            price_array_here = np.divide(1, 1 - pcm_array)
             price_array = price_array_here[:, :2]
-            nth_firm_price = price_array_here[:, [-1]]  # noqa: F841
+            nth_firm_price = price_array_here[:, [-1]]
             if _pcm_spec.pcm_restriction == PCMRestriction.MNL:
                 # Generate i.i.d. PCMs then take PCM0 and construct PCM1
                 # Regenerate MNL test
-                purchase_prob_array = ne.evaluate(
-                    "_aggregate_purchase_prob * _frmshr_array"
+                purchase_prob_array = np.einsum(
+                    "ij,ij->ij", _aggregate_purchase_prob, _frmshr_array
                 )
 
-                _pr_0, _pi_0 = price_array_here[:, [0]], purchase_prob_array[:, [0]]
-                _pcm_0 = pcm_array[:, [0]]
-                _pr_1, _pi_1 = price_array_here[:, [1]], purchase_prob_array[:, [1]]
-
-                pcm_array[:, [1]] = ne.evaluate(
-                    "_pcm_0 * (1 - _pi_0) / (1 - _pi_1 + _pcm_0 * (_pi_1 - _pi_0))"
+                pcm_array[:, 1] = np.divide(
+                    np.einsum("i,i->i", pcm_array[:, 0], 1 - purchase_prob_array[:, 0]),
+                    1
+                    - purchase_prob_array[:, 1]
+                    + np.einsum(
+                        "i,i->i",
+                        pcm_array[:, 0],
+                        purchase_prob_array[:, 1] - purchase_prob_array[:, 0],
+                    ),
                 )
-                mnl_test_array = (pcm_array[:, [1]] >= 0) & (pcm_array[:, [1]] <= 1)
 
-            margin_data = MarginSampleData(pcm_array[:, :2], mnl_test_array)
+                mnl_test = (pcm_array[:, [1]] >= 0) & (pcm_array[:, [1]] <= 1)
+
+            margin_data = MarginsData(pcm_array[:, :2], mnl_test)
             del price_array_here
         case _:
             raise ValueError(
                 f'Specification of price distribution, "{_price_spec.value}" is invalid.'
             )
     if _price_spec != PriceSpec.CSY:
-        margin_data = _gen_margin_data(
+        margin_data = _margins_sampler(
             _frmshr_array,
             price_array,
             _aggregate_purchase_prob,
@@ -600,8 +605,8 @@ def gen_margin_price_data(
         )
 
     # _price_array = _price_array.astype(np.float64)
-    rev_array = ne.evaluate("price_array * _frmshr_array")
-    nth_firm_rev = ne.evaluate("nth_firm_price * _nth_firm_share")
+    rev_array = np.einsum("ij,ij->ij", price_array, _frmshr_array)
+    nth_firm_rev = np.einsum("ij,ij->ij", nth_firm_price, _nth_firm_share)
 
     # Although `_test_rev_ratio_inv` is not fixed at 10%,
     # the ratio has not changed since inception of the HSR filing test,
@@ -646,10 +651,10 @@ def gen_margin_price_data(
     # under value-of-transactions test or merger triggers post-consummation review
     hsr_filing_test |= _frmshr_array.min(axis=1) >= test_rev_ratio_inv
 
-    return margin_data, PriceSampleData(price_array, hsr_filing_test)
+    return margin_data, PricesData(price_array, hsr_filing_test)
 
 
-def _gen_margin_data(
+def _margins_sampler(
     _frmshr_array: ArrayDouble,
     _price_array: ArrayDouble,
     _aggregate_purchase_prob: ArrayDouble,
@@ -657,7 +662,7 @@ def _gen_margin_data(
     _pcm_rng_seed_seq: SeedSequence,
     _nthreads: int,
     /,
-) -> MarginSampleData:
+) -> MarginsData:
     dist_type_pcm, dist_parms_pcm, pcm_restriction_ = (
         getattr(_pcm_spec, _f) for _f in ("dist_type", "dist_parms", "pcm_restriction")
     )
@@ -722,19 +727,20 @@ def _gen_margin_data(
         # Impose FOCs from profit-maximization with MNL demand
         purchase_prob_array = _aggregate_purchase_prob * _frmshr_array
 
-        _pr_0, _pi_0 = _price_array[:, [0]], purchase_prob_array[:, [0]]
-        _pcm_0 = pcm_array[:, [0]]
-        _pr_1, _pi_1 = _price_array[:, [1]], purchase_prob_array[:, [1]]
-
-        pcm_array[:, [1]] = ne.evaluate(
-            "_pr_0 * _pcm_0 * (1 - _pi_0) / (_pr_1 * (1 - _pi_1))"
+        pcm_array[:, 1] = np.divide(
+            np.einsum(
+                "i,i,i->i",
+                _price_array[:, 0],
+                pcm_array[:, 0],
+                1 - purchase_prob_array[:, 0],
+            ),
+            np.einsum("i,i->i", _price_array[:, 1], 1 - purchase_prob_array[:, 1]),
         )
-
-        mnl_test_array = (pcm_array[:, 1] >= 0) & (pcm_array[:, 1] <= 1)
+        mnl_test = (pcm_array[:, 1] >= 0) & (pcm_array[:, 1] <= 1)
     else:
-        mnl_test_array = np.ones(len(pcm_array), dtype=bool)
+        mnl_test = np.ones(len(pcm_array), dtype=bool)
 
-    return MarginSampleData(pcm_array, mnl_test_array)
+    return MarginsData(pcm_array, mnl_test)
 
 
 def _beta_located(

mergeron/gen/enforcement_stats.py

@@ -335,12 +335,12 @@ def table_no_lku(
     return tno_
 
 
-def enf_cnts_byfirmcount(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
+def enf_cnts_byfirmcount(_raw_counts: ArrayBIGINT, /) -> ArrayBIGINT:
     """Summarize investigations data by firm count.
 
     Parameters
     ----------
-    _cnts_array
+    _raw_counts
         raw investigations data array
 
     Returns
@@ -348,7 +348,7 @@ def enf_cnts_byfirmcount(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
     ArrayBIGINT
         Subtotals for columns other than the first, grouped by the first column.
     """
-    if not _cnts_array[:, 0].any():
+    if not _raw_counts[:, 0].any():
         return np.array([], int)
 
     ndim_in = 1
@@ -356,19 +356,19 @@ def enf_cnts_byfirmcount(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
         np.concatenate([
             (_i,),
             np.einsum(
-                "ij->j", _cnts_array[_cnts_array[:, 0] == _i][:, ndim_in:], dtype=int
+                "ij->j", _raw_counts[_raw_counts[:, 0] == _i][:, ndim_in:], dtype=int
             ),
         ])
-        for _i in np.unique(_cnts_array[:, 0])
+        for _i in np.unique(_raw_counts[:, 0])
     ])
 
 
-def enf_cnts_bydelta(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
+def enf_cnts_bydelta(_raw_counts: ArrayBIGINT, /) -> ArrayBIGINT:
     """Summarize investigations data by ΔHHI.
 
     Parameters
     ----------
-    _cnts_array
+    _raw_counts
         raw investigations data array
 
     Returns
@@ -376,28 +376,28 @@ def enf_cnts_bydelta(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
     ArrayBIGINT
         Subtotals for columns higher than the second, grouped by the second column.
     """
-    if not _cnts_array[:, 1].any():
+    if not _raw_counts[:, 1].any():
         return np.array([], int)
     ndim_in = 2
     return np.vstack([
         np.concatenate([
             (_k,),
             np.einsum(
-                "ij->j", _cnts_array[_cnts_array[:, 1] == _k][:, ndim_in:], dtype=int
+                "ij->j", _raw_counts[_raw_counts[:, 1] == _k][:, ndim_in:], dtype=int
             ),
         ])
         for _k in HHI_DELTA_KNOTS[:-1]
     ])
 
 
-def enf_cnts_byconczone(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
+def enf_cnts_byconczone(_raw_counts: ArrayBIGINT, /) -> ArrayBIGINT:
     """Summarize investigations data by concentration zone, as defined in the Guidelines.
 
     Includes sub-total detail for "Moderately Concentrated" and "Unconcentrated" markets.
 
     Parameters
     ----------
-    _cnts_array
+    _raw_counts
         raw investigations data array
 
     Returns
@@ -405,7 +405,7 @@ def enf_cnts_byconczone(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
     ArrayBIGINT
         Subtotals range of HHI and ΔHHI, with detail
     """
-    if not _cnts_array[:, 0].any() or np.isnan(_cnts_array[:, 0]).all():
+    if not _raw_counts[:, 0].any() or np.isnan(_raw_counts[:, 0]).all():
         return np.array([], int)
     # Step 1: Tag and agg. from HHI-post and ΔHHI to zone triple
     # NOTE: Although you could just map and not (partially) aggregate in this step,
@@ -416,15 +416,15 @@ def enf_cnts_byconczone(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
     _ndim_in = 2
     _nkeys = 3
     cnts_byhhipostanddelta: ArrayBIGINT = np.zeros(
-        (1, _nkeys + _cnts_array.shape[1] - _ndim_in), dtype=int
+        (1, _nkeys + _raw_counts.shape[1] - _ndim_in), dtype=int
     )
     cnts_byconczone: ArrayBIGINT = np.zeros(
-        (1, _nkeys + _cnts_array.shape[1] - _ndim_in), dtype=int
+        (1, _nkeys + _raw_counts.shape[1] - _ndim_in), dtype=int
     )
 
     # Prepare to tag clearance stats by presumption zone
-    hhi_zone_post_ranged = hhi_zone_post_ranger(_cnts_array[:, 0] / 1e4)
-    hhi_delta_ranged = hhi_delta_ranger(_cnts_array[:, 1] / 1e4)
+    hhi_zone_post_ranged = hhi_zone_post_ranger(_raw_counts[:, 0] / 1e4)
+    hhi_delta_ranged = hhi_delta_ranger(_raw_counts[:, 1] / 1e4)
     for _hhi_zone_post_lim in HHI_POST_ZONE_KNOTS[:-1]:
         zone_test = hhi_zone_post_ranged == _hhi_zone_post_lim
 
@@ -445,7 +445,7 @@ def enf_cnts_byconczone(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
                     (
                         *zone_val,
                         *np.einsum(
-                            "ij->j", _cnts_array[:, _ndim_in:][conc_test], dtype=int
+                            "ij->j", _raw_counts[:, _ndim_in:][conc_test], dtype=int
                         ),
                     ),
                     dtype=int,