mergeron 2025.739290.3__py3-none-any.whl → 2025.739290.4__py3-none-any.whl

mergeron/gen/enforcement_stats.py

@@ -9,15 +9,16 @@ from collections.abc import Mapping
 import numpy as np
 from scipy.interpolate import interp1d  # type: ignore
 
-from .. import VERSION, ArrayBIGINT, this_yaml  # noqa: TID252
+from .. import VERSION, ArrayBIGINT, EnumYAMLized, this_yaml  # noqa: TID252
 from ..core import ftc_merger_investigations_data as fid  # noqa: TID252
 from . import INVResolution
 
 __version__ = VERSION
 
 
+@this_yaml.register_class
 @enum.unique
-class IndustryGroup(enum.StrEnum):
+class IndustryGroup(str, EnumYAMLized):
     ALL = "All Markets"
     GRO = "Grocery Markets"
     OIL = "Oil Markets"
@@ -30,8 +31,9 @@ class IndustryGroup(enum.StrEnum):
     IIC = "Industries in Common"
 
 
+@this_yaml.register_class
 @enum.unique
-class OtherEvidence(enum.StrEnum):
+class OtherEvidence(str, EnumYAMLized):
     UR = "Unrestricted on additional evidence"
     HD = "Hot Documents Identified"
     HN = "No Hot Documents Identified"
@@ -44,23 +46,26 @@ class OtherEvidence(enum.StrEnum):
     NE = "No Entry Evidence"
 
 
+@this_yaml.register_class
 @enum.unique
-class StatsGrpSelector(enum.StrEnum):
+class StatsGrpSelector(str, EnumYAMLized):
     FC = "ByFirmCount"
     HD = "ByHHIandDelta"
     DL = "ByDelta"
     ZN = "ByConcZone"
 
 
+@this_yaml.register_class
 @enum.unique
-class StatsReturnSelector(enum.StrEnum):
+class StatsReturnSelector(str, EnumYAMLized):
     CNT = "count"
     RPT = "rate, point"
     RIN = "rate, interval"
 
 
+@this_yaml.register_class
 @enum.unique
-class SortSelector(enum.StrEnum):
+class SortSelector(str, EnumYAMLized):
     UCH = "unchanged"
     REV = "reversed"
 
@@ -138,17 +143,17 @@ def enf_cnts_obs_by_group(
 
     match _stats_group:
         case StatsGrpSelector.FC:
-            _cnts_func = enf_cnts_byfirmcount
-            _cnts_listing_func = enf_cnts_obs_byfirmcount
+            cnts_func = enf_cnts_byfirmcount
+            cnts_listing_func = enf_cnts_obs_byfirmcount
         case StatsGrpSelector.DL:
-            _cnts_func = enf_cnts_bydelta
-            _cnts_listing_func = enf_cnts_obs_byhhianddelta
+            cnts_func = enf_cnts_bydelta
+            cnts_listing_func = enf_cnts_obs_byhhianddelta
         case StatsGrpSelector.ZN:
-            _cnts_func = enf_cnts_byconczone
-            _cnts_listing_func = enf_cnts_obs_byhhianddelta
+            cnts_func = enf_cnts_byconczone
+            cnts_listing_func = enf_cnts_obs_byhhianddelta
 
-    return _cnts_func(
-        _cnts_listing_func(
+    return cnts_func(
+        cnts_listing_func(
             _invdata_array_dict,
             _study_period,
             _table_ind_grp,
@@ -172,26 +177,23 @@ def enf_cnts_obs_byfirmcount(
             f"Must be one of, {tuple(_data_array_dict.keys())!r}."
         )
 
-    _data_array_dict_sub = _data_array_dict[_data_period][fid.TABLE_TYPES[1]]
+    data_array_dict_sub = _data_array_dict[_data_period][fid.TABLE_TYPES[1]]
 
-    _table_no = table_no_lku(_data_array_dict_sub, _table_ind_group, _table_evid_cond)
+    table_no_ = table_no_lku(data_array_dict_sub, _table_ind_group, _table_evid_cond)
 
-    _cnts_array = _data_array_dict_sub[_table_no].data_array
+    cnts_array = data_array_dict_sub[table_no_].data_array
 
-    _ndim_in = 1
-    _stats_kept_indxs = []
+    ndim_in = 1
+    stats_kept_indxs = []
     match _enf_spec:
         case INVResolution.CLRN:
-            _stats_kept_indxs = [-1, -2]
+            stats_kept_indxs = [-1, -2]
         case INVResolution.ENFT:
-            _stats_kept_indxs = [-1, -3]
+            stats_kept_indxs = [-1, -3]
         case INVResolution.BOTH:
-            _stats_kept_indxs = [-1, -3, -2]
+            stats_kept_indxs = [-1, -3, -2]
 
-    return np.column_stack([
-        _cnts_array[:, :_ndim_in],
-        _cnts_array[:, _stats_kept_indxs],
-    ])
+    return np.column_stack([cnts_array[:, :ndim_in], cnts_array[:, stats_kept_indxs]])
 
 
 def enf_cnts_obs_byhhianddelta(
@@ -208,26 +210,23 @@ def enf_cnts_obs_byhhianddelta(
             f"Must be one of, {tuple(_data_array_dict.keys())!r}."
         )
 
-    _data_array_dict_sub = _data_array_dict[_data_period][fid.TABLE_TYPES[0]]
+    data_array_dict_sub = _data_array_dict[_data_period][fid.TABLE_TYPES[0]]
 
-    _table_no = table_no_lku(_data_array_dict_sub, _table_ind_group, _table_evid_cond)
+    table_no_ = table_no_lku(data_array_dict_sub, _table_ind_group, _table_evid_cond)
 
-    _cnts_array = _data_array_dict_sub[_table_no].data_array
+    cnts_array = data_array_dict_sub[table_no_].data_array
 
-    _ndim_in = 2
-    _stats_kept_indxs = []
+    ndim_in = 2
+    stats_kept_indxs = []
     match _enf_spec:
         case INVResolution.CLRN:
-            _stats_kept_indxs = [-1, -2]
+            stats_kept_indxs = [-1, -2]
         case INVResolution.ENFT:
-            _stats_kept_indxs = [-1, -3]
+            stats_kept_indxs = [-1, -3]
         case INVResolution.BOTH:
-            _stats_kept_indxs = [-1, -3, -2]
+            stats_kept_indxs = [-1, -3, -2]
 
-    return np.column_stack([
-        _cnts_array[:, :_ndim_in],
-        _cnts_array[:, _stats_kept_indxs],
-    ])
+    return np.column_stack([cnts_array[:, :ndim_in], cnts_array[:, stats_kept_indxs]])
 
 
 def table_no_lku(
@@ -237,143 +236,112 @@ def table_no_lku(
     /,
 ) -> str:
     if _table_ind_group not in (
-        _igl := [_data_array_dict_sub[_v].industry_group for _v in _data_array_dict_sub]
+        igl_ := [_data_array_dict_sub[_v].industry_group for _v in _data_array_dict_sub]
     ):
         raise ValueError(
             f"Invalid value for industry group, {f'"{_table_ind_group}"'}."
-            f"Must be one of {_igl!r}"
+            f"Must be one of {igl_!r}"
         )
 
-    _tno = next(
-        _t
-        for _t in _data_array_dict_sub
+    tno_ = next(
+        t_
+        for t_ in _data_array_dict_sub
         if all((
-            _data_array_dict_sub[_t].industry_group == _table_ind_group,
-            _data_array_dict_sub[_t].additional_evidence == _table_evid_cond,
+            _data_array_dict_sub[t_].industry_group == _table_ind_group,
+            _data_array_dict_sub[t_].additional_evidence == _table_evid_cond,
         ))
     )
 
-    return _tno
+    return tno_
 
 
 def enf_cnts_byfirmcount(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
-    _ndim_in = 1
+    ndim_in = 1
     return np.vstack([
         np.concatenate([
             (f,),
-            np.einsum("ij->j", _cnts_array[_cnts_array[:, 0] == f][:, _ndim_in:]),
+            np.einsum("ij->j", _cnts_array[_cnts_array[:, 0] == f][:, ndim_in:]),
         ])
         for f in np.unique(_cnts_array[:, 0])
     ])
 
 
 def enf_cnts_bydelta(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
-    _ndim_in = 2
+    ndim_in = 2
     return np.vstack([
         np.concatenate([
-            (f,),
-            np.einsum("ij->j", _cnts_array[_cnts_array[:, 1] == f][:, _ndim_in:]),
+            (f_,),
+            np.einsum("ij->j", _cnts_array[_cnts_array[:, 1] == f_][:, ndim_in:]),
         ])
-        for f in HHI_DELTA_KNOTS[:-1]
+        for f_ in HHI_DELTA_KNOTS[:-1]
     ])
 
 
 def enf_cnts_byconczone(_cnts_array: ArrayBIGINT, /) -> ArrayBIGINT:
-    # 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)
-
+    if not _cnts_array.any():
+        return np.array([], int)
     # Step 1: Tag and agg. from HHI-post and Delta to zone triple
     # NOTE: Although you could just map and not (partially) aggregate in this step,
     # the mapped array is a copy, and is larger without partial aggregation, so
     # aggregation reduces the footprint of this step in memory. Although this point
     # is more relevant for generated than observed data, using the same coding pattern
     # in both cases does make life easier
-    _ndim_in = 2
-    _nkeys = 3
-    _cnts_byhhipostanddelta: ArrayBIGINT = -1 * np.ones(
-        _nkeys + _cnts_array.shape[1] - _ndim_in, dtype=int
+    ndim_in = 2
+    nkeys_ = 3
+    cnts_byhhipostanddelta, cnts_byconczone = (
+        np.zeros(nkeys_ + _cnts_array.shape[1] - ndim_in, dtype=int) for _ in range(2)
     )
-    _cnts_byconczone: ArrayBIGINT = -1 * np.ones_like(_cnts_byhhipostanddelta)
+
+    # 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)
     for _hhi_zone_post_lim in HHI_POST_ZONE_KNOTS[:-1]:
-        _level_test = _hhi_zone_post_ranged == _hhi_zone_post_lim
+        zone_test = hhi_zone_post_ranged == _hhi_zone_post_lim
 
-        for _hhi_zone_delta_lim in HHI_DELTA_KNOTS[:3]:
-            _delta_test = (
-                (_hhi_delta_ranged >= _hhi_zone_delta_lim)
-                if _hhi_zone_delta_lim == HHI_DELTA_KNOTS[2]
-                else (_hhi_delta_ranged == _hhi_zone_delta_lim)
+        for hhi_zone_delta_lim in HHI_DELTA_KNOTS[:3]:
+            delta_test = (
+                (hhi_delta_ranged >= hhi_zone_delta_lim)
+                if hhi_zone_delta_lim == HHI_DELTA_KNOTS[2]
+                else (hhi_delta_ranged == hhi_zone_delta_lim)
             )
 
-            _zone_val = HMG_PRESUMPTION_ZONE_MAP[_hhi_zone_post_lim][
-                _hhi_zone_delta_lim
-            ]
+            zone_val = HMG_PRESUMPTION_ZONE_MAP[_hhi_zone_post_lim][hhi_zone_delta_lim]
 
-            _conc_test = _level_test & _delta_test
+            conc_test = zone_test & delta_test
 
-            _cnts_byhhipostanddelta = np.vstack((
-                _cnts_byhhipostanddelta,
+            cnts_byhhipostanddelta = np.vstack((
+                cnts_byhhipostanddelta,
                 np.array(
                     (
-                        *_zone_val,
-                        *np.einsum("ij->j", _cnts_array[:, _ndim_in:][_conc_test]),
+                        *zone_val,
+                        *np.einsum("ij->j", _cnts_array[:, ndim_in:][conc_test]),
                     ),
                     dtype=int,
                 ),
             ))
-    _cnts_byhhipostanddelta = _cnts_byhhipostanddelta[1:]
+    cnts_byhhipostanddelta = cnts_byhhipostanddelta[1:]
 
-    for _zone_val in ZONE_VALS:
+    for zone_val in ZONE_VALS:
         # Logical-and of multiple vectors:
-        _hhi_zone_test = (
+        hhi_zone_test = (
             1
             * np.column_stack([
-                _cnts_byhhipostanddelta[:, _idx] == _val
-                for _idx, _val in enumerate(_zone_val)
+                cnts_byhhipostanddelta[:, _idx] == _val
+                for _idx, _val in enumerate(zone_val)
             ])
         ).prod(axis=1) == 1
 
-        _cnts_byconczone = np.vstack((
-            _cnts_byconczone,
+        cnts_byconczone = np.vstack((
+            cnts_byconczone,
             np.concatenate(
                 (
-                    _zone_val,
+                    zone_val,
                     np.einsum(
-                        "ij->j", _cnts_byhhipostanddelta[_hhi_zone_test][:, _nkeys:]
+                        "ij->j", cnts_byhhipostanddelta[hhi_zone_test][:, nkeys_:]
                     ),
                 ),
                 dtype=int,
             ),
         ))
 
-    return _cnts_byconczone[1:]
-
-
-for _typ in (
-    IndustryGroup,
-    OtherEvidence,
-    StatsGrpSelector,
-    StatsReturnSelector,
-    SortSelector,
-):
-    #  NOTE: If additional enums are defined in this module,
-    #  add themn to the list above
-
-    _, _ = (
-        this_yaml.representer.add_representer(
-            _typ,
-            lambda _r, _d: _r.represent_scalar(f"!{_d.__class__.__name__}", _d.name),
-        ),
-        this_yaml.constructor.add_constructor(
-            f"!{_typ.__name__}",
-            lambda _c, _n, /: getattr(
-                globals().get(_n.tag.lstrip("!")), _c.construct_scalar(_n)
-            ),
-        ),
-    )
-
-
-if __name__ == "__main__":
-    print(
-        "This module provides methods to aggregate statistics on merger enforcement patterns for reporting."
-    )
+    return cnts_byconczone[1:]