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

mergeron/gen/__init__.py

@@ -8,14 +8,11 @@ from __future__ import annotations
 
 import enum
 from collections.abc import Sequence
-from dataclasses import dataclass
-from typing import ClassVar, Protocol
+from operator import attrgetter
 
 import numpy as np
-from attrs import Attribute, Converter, cmp_using, field, validators
-from attrs import frozen as frozen_attrs
+from attrs import Attribute, Converter, cmp_using, field, frozen, validators
 from numpy.random import SeedSequence
-from ruamel import yaml
 
 from .. import (  # noqa: TID252
     DEFAULT_REC_RATIO,
@@ -25,9 +22,11 @@ from .. import (  # noqa: TID252
     ArrayDouble,
     ArrayFloat,
     ArrayINT,
+    EnumYAMLized,
     RECForm,
     UPPAggrSelector,
     this_yaml,
+    yamelize_attrs,
 )
 from ..core.pseudorandom_numbers import (  # noqa: TID252
     DEFAULT_BETA_DIST_PARMS,
@@ -36,28 +35,29 @@ from ..core.pseudorandom_numbers import (  # noqa: TID252
 
 __version__ = VERSION
 
-
 DEFAULT_FCOUNT_WTS = np.asarray((_nr := np.arange(6, 0, -1)) / _nr.sum(), float)
 
+DEFAULT_BETA_BND_DIST_PARMS = np.array([0.5, 1.0, 0.0, 1.0], float)
+
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class SeedSequenceData:
-    share: SeedSequence
-    pcm: SeedSequence
-    fcounts: SeedSequence | None
-    price: SeedSequence | None
+    share: SeedSequence = field(eq=attrgetter("state"))
+    pcm: SeedSequence = field(eq=attrgetter("state"))
+    fcounts: SeedSequence | None = field(eq=lambda x: x if x is None else x.state)
+    price: SeedSequence | None = field(eq=lambda x: x if x is None else x.state)
 
 
 @this_yaml.register_class
 @enum.unique
-class PriceSpec(tuple[bool, str | None], enum.ReprEnum):
+class PriceSpec(tuple[bool, str | None], EnumYAMLized):
     """Price specification.
 
     Whether prices are symmetric and, if not, the direction of correlation, if any.
     """
 
     SYM = (True, None)
-    ZERO = (False, None)
+    RNG = (False, None)
     NEG = (False, "negative share-correlation")
     POS = (False, "positive share-correlation")
     CSY = (False, "market-wide cost-symmetry")
@@ -65,7 +65,7 @@ class PriceSpec(tuple[bool, str | None], enum.ReprEnum):
 
 @this_yaml.register_class
 @enum.unique
-class SHRDistribution(enum.StrEnum):
+class SHRDistribution(str, EnumYAMLized):
     """Market share distributions."""
 
     UNI = "Uniform"
@@ -104,30 +104,28 @@ def _fc_wts_conv(
 ) -> ArrayFloat | None:
     if _i.dist_type == SHRDistribution.UNI:
         return None
-    elif _v is None or np.array_equal(_v, DEFAULT_FCOUNT_WTS):
+    elif _v is None or len(_v) == 0 or np.array_equal(_v, DEFAULT_FCOUNT_WTS):
         return DEFAULT_FCOUNT_WTS
     else:
         return _tv if (_tv := np.asarray(_v, float)).sum() == 1 else _tv / _tv.sum()
 
 
 def _shr_dp_conv(_v: Sequence[float] | ArrayFloat | None, _i: ShareSpec) -> ArrayFloat:
-    match _v:
-        case None if _i.dist_type == SHRDistribution.UNI:
+    if _v is None or len(_v) == 0 or np.array_equal(_v, DEFAULT_DIST_PARMS):
+        if _i.dist_type == SHRDistribution.UNI:
             return DEFAULT_DIST_PARMS
-        case None:
-            _fc_max = 1 + (
+        else:
+            fc_max = 1 + (
                 len(DEFAULT_FCOUNT_WTS)
-                if not hasattr(_i, "firm_counts_weights")
-                or _i.firm_counts_weights is None
-                or len(_i.firm_counts_weights) == 0
+                if _i.firm_counts_weights is None
                 else len(_i.firm_counts_weights)
             )
 
             match _i.dist_type:
                 case SHRDistribution.DIR_FLAT | SHRDistribution.DIR_FLAT_CONSTR:
-                    return np.ones(_fc_max, float)
+                    return np.ones(fc_max, float)
                 case SHRDistribution.DIR_ASYM:
-                    return np.array([2.0] * 6 + [1.5] * 5 + [1.25] * _fc_max, float)
+                    return np.array([2.0] * 6 + [1.5] * 5 + [1.25] * fc_max, float)
                 case SHRDistribution.DIR_COND:
                     return np.array([], float)
                 case _ if isinstance(_i.dist_type, SHRDistribution):
@@ -138,16 +136,15 @@ def _shr_dp_conv(_v: Sequence[float] | ArrayFloat | None, _i: ShareSpec) -> Arra
                     raise ValueError(
                         f"Unsupported distribution for market share generation, {_i.dist_type!r}"
                     )
-        case _ if isinstance(_v, Sequence | np.ndarray):
-            return np.asarray(_v, float)
-        case _:
-            raise ValueError(
-                f"Input, {_v!r} has invalid type. Must be None, Sequence of floats, or Numpy ndarray."
-            )
+    elif isinstance(_v, Sequence | np.ndarray):
+        return np.asarray(_v, float)
+    else:
+        raise ValueError(
+            f"Input, {_v!r} has invalid type. Must be None, Sequence of floats, or Numpy ndarray."
+        )
 
 
-@this_yaml.register_class
-@frozen_attrs
+@frozen
 class ShareSpec:
     """Market share specification
 
@@ -188,11 +185,11 @@ class ShareSpec:
     """
 
     @firm_counts_weights.default
-    def _fcwd(_i: ShareSpec) -> ArrayFloat | None:
+    def __fcwd(_i: ShareSpec) -> ArrayFloat | None:
         return _fc_wts_conv(None, _i)
 
     @firm_counts_weights.validator
-    def _fcv(_i: ShareSpec, _a: Attribute[ArrayFloat], _v: ArrayFloat) -> None:
+    def __fcv(_i: ShareSpec, _a: Attribute[ArrayFloat], _v: ArrayFloat) -> None:
         if _i.dist_type != SHRDistribution.UNI and not len(_v):
             raise ValueError(
                 f"Attribute, {'"firm_counts_weights"'} must be populated if the share distribution is "
@@ -214,13 +211,13 @@ class ShareSpec:
     """
 
     @dist_parms.default
-    def _dpd(_i: ShareSpec) -> ArrayFloat:
+    def __dpd(_i: ShareSpec) -> ArrayFloat:
         # converters run after defaults, and we
         # avoid redundancy and confusion here
         return _shr_dp_conv(None, _i)
 
     @dist_parms.validator
-    def _dpv(_i: ShareSpec, _a: Attribute[ArrayFloat], _v: ArrayFloat) -> None:
+    def __dpv(_i: ShareSpec, _a: Attribute[ArrayFloat], _v: ArrayFloat) -> None:
         if (
             _i.firm_counts_weights is not None
             and _v is not None
@@ -236,7 +233,7 @@ class ShareSpec:
     """See :class:`mergeron.RECForm`"""
 
     @recapture_form.validator
-    def _rfv(_i: ShareSpec, _a: Attribute[RECForm], _v: RECForm) -> None:
+    def __rfv(_i: ShareSpec, _a: Attribute[RECForm], _v: RECForm) -> None:
         if _i.dist_type == SHRDistribution.UNI and _v == RECForm.OUTIN:
             raise ValueError(
                 "Outside-good choice probabilities cannot be generated if the "
@@ -271,11 +268,11 @@ class ShareSpec:
     """
 
     @recapture_ratio.default
-    def _rrd(_i: ShareSpec) -> float | None:
+    def __rrd(_i: ShareSpec) -> float | None:
         return None if _i.recapture_form == RECForm.OUTIN else DEFAULT_REC_RATIO
 
     @recapture_ratio.validator
-    def _rrv(_i: ShareSpec, _a: Attribute[float], _v: float) -> None:
+    def __rrv(_i: ShareSpec, _a: Attribute[float], _v: float) -> None:
         if _v and not (0 < _v <= 1):
             raise ValueError("Recapture ratio must lie in the interval, [0, 1).")
         elif _v is None and _i.recapture_form != RECForm.OUTIN:
@@ -285,25 +282,10 @@ class ShareSpec:
                 "interval [0, 1)."
             )
 
-    @classmethod
-    def to_yaml(
-        cls, _r: yaml.representer.SafeRepresenter, _d: ShareSpec
-    ) -> yaml.MappingNode:
-        _ret: yaml.MappingNode = _r.represent_mapping(
-            f"!{cls.__name__}",
-            {_a.name: getattr(_d, _a.name) for _a in _d.__attrs_attrs__},
-        )
-        return _ret
-
-    @classmethod
-    def from_yaml(
-        cls, _c: yaml.constructor.SafeConstructor, _n: yaml.MappingNode
-    ) -> ShareSpec:
-        return cls(**_c.construct_mapping(_n))
-
 
+@this_yaml.register_class
 @enum.unique
-class PCMDistribution(enum.StrEnum):
+class PCMDistribution(str, EnumYAMLized):
     """Margin distributions."""
 
     UNI = "Uniform"
@@ -312,8 +294,9 @@ class PCMDistribution(enum.StrEnum):
     EMPR = "Damodaran margin data, resampled"
 
 
+@this_yaml.register_class
 @enum.unique
-class FM2Constraint(enum.StrEnum):
+class FM2Constraint(str, EnumYAMLized):
     """Firm 2 margins - derivation methods."""
 
     IID = "i.i.d"
@@ -321,8 +304,28 @@ class FM2Constraint(enum.StrEnum):
     SYM = "symmetric"
 
 
-@this_yaml.register_class
-@frozen_attrs
+def _pcm_dp_conv(
+    _v: Sequence[float] | ArrayFloat | None, _i: PCMSpec
+) -> ArrayFloat | None:
+    if _i.dist_type == PCMDistribution.EMPR:
+        return None
+    elif _v is None or len(_v) == 0 or np.array_equal(_v, DEFAULT_DIST_PARMS):
+        match _i.dist_type:
+            case PCMDistribution.BETA:
+                return DEFAULT_BETA_DIST_PARMS
+            case PCMDistribution.BETA_BND:
+                return DEFAULT_BETA_BND_DIST_PARMS
+            case _:
+                return DEFAULT_DIST_PARMS
+    elif isinstance(_v, Sequence | np.ndarray):
+        return np.asarray(_v, float)
+    else:
+        raise ValueError(
+            f"Input, {_v!r} has invalid type. Must be None, Sequence of floats, or Numpy ndarray."
+        )
+
+
+@frozen
 class PCMSpec:
     """Price-cost margin (PCM) specification
 
@@ -342,10 +345,14 @@ class PCMSpec:
     """See :class:`PCMDistribution`"""
 
     @dist_type.default
-    def _dtd(_i: PCMSpec) -> PCMDistribution:
+    def __dtd(_i: PCMSpec) -> PCMDistribution:
         return PCMDistribution.UNI
 
-    dist_parms: ArrayFloat | None = field(kw_only=True, eq=np.array_repr)
+    dist_parms: ArrayFloat | None = field(
+        kw_only=True,
+        eq=cmp_using(eq=np.array_equal),
+        converter=Converter(_pcm_dp_conv, takes_self=True),  # type: ignore
+    )
     """Parameter specification for tailoring PCM distribution
 
     For Uniform distribution, bounds of the distribution; defaults to `(0, 1)`;
@@ -356,19 +363,11 @@ class PCMSpec:
     """
 
     @dist_parms.default
-    def _dpwd(_i: PCMSpec) -> ArrayFloat | None:
-        match _i.dist_type:
-            case PCMDistribution.EMPR:
-                return None
-            case PCMDistribution.BETA:
-                return DEFAULT_BETA_DIST_PARMS
-            case PCMDistribution.BETA_BND:
-                return np.array([0.5, 1.0, 0.0, 0.1], float)
-            case _:
-                return DEFAULT_DIST_PARMS
+    def __dpwd(_i: PCMSpec) -> ArrayFloat | None:
+        return _pcm_dp_conv(None, _i)
 
     @dist_parms.validator
-    def _dpv(
+    def __dpv(
         _i: PCMSpec, _a: Attribute[ArrayFloat | None], _v: ArrayFloat | None
     ) -> None:
         if _i.dist_type.name.startswith("BETA"):
@@ -399,25 +398,10 @@ class PCMSpec:
     firm2_pcm_constraint: FM2Constraint = field(kw_only=True, default=FM2Constraint.IID)
     """See :class:`FM2Constraint`"""
 
-    @classmethod
-    def to_yaml(
-        cls, _r: yaml.representer.SafeRepresenter, _d: PCMSpec
-    ) -> yaml.MappingNode:
-        _ret: yaml.MappingNode = _r.represent_mapping(
-            f"!{cls.__name__}",
-            {_a.name: getattr(_d, _a.name) for _a in _d.__attrs_attrs__},
-        )
-        return _ret
-
-    @classmethod
-    def from_yaml(
-        cls, _c: yaml.constructor.SafeConstructor, _n: yaml.MappingNode
-    ) -> PCMSpec:
-        return cls(**_c.construct_mapping(_n))
-
 
+@this_yaml.register_class
 @enum.unique
-class SSZConstant(float, enum.ReprEnum):
+class SSZConstant(float, EnumYAMLized):
     """
     Scale factors to offset sample size reduction.
 
@@ -456,47 +440,69 @@ class SSZConstant(float, enum.ReprEnum):
     """When initial set of draws is not restricted in any way."""
 
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class MarketSampleData:
     """Container for generated markets data sample."""
 
-    frmshr_array: ArrayDouble
+    frmshr_array: ArrayDouble = field(eq=cmp_using(np.array_equal))
     """Merging-firm shares (with two merging firms)"""
 
-    pcm_array: ArrayDouble
+    pcm_array: ArrayDouble = field(eq=cmp_using(np.array_equal))
     """Merging-firms' prices (normalized to 1, in default specification)"""
 
-    price_array: ArrayDouble
+    price_array: ArrayDouble = field(eq=cmp_using(np.array_equal))
     """Merging-firms' price-cost margins (PCM)"""
 
-    fcounts: ArrayBIGINT
-    """Number of firms in market"""
+    divr_array: ArrayDouble = field(eq=cmp_using(np.array_equal))
+    """Diversion ratio between the merging firms"""
 
-    aggregate_purchase_prob: ArrayDouble
+    hhi_delta: ArrayDouble = field(eq=cmp_using(np.array_equal))
+    """Change in HHI from combination of merging firms"""
+
+    aggregate_purchase_prob: ArrayDouble = field(eq=cmp_using(np.array_equal))
     """
     One (1) minus probability that the outside good is chosen
 
     Converts market shares to choice probabilities by multiplication.
     """
 
-    nth_firm_share: ArrayDouble
+    @aggregate_purchase_prob.default
+    def __appd(_i: MarketSampleData) -> ArrayINT:
+        e_ = np.empty_like(_i.frmshr_array[:, :1], float)
+        e_.fill(np.nan)
+        return e_
+
+    fcounts: ArrayINT = field(eq=cmp_using(np.array_equal))
+    """Number of firms in market"""
+
+    @fcounts.default
+    def __fcd(_i: MarketSampleData) -> ArrayINT:
+        return np.zeros_like(_i.frmshr_array[:, :1], np.uint8)
+
+    nth_firm_share: ArrayDouble = field(eq=cmp_using(np.array_equal))
     """Market-share of n-th firm
 
-    Relevant for testing for draws the do or
+    Relevant for testing draws that do or
     do not meet HSR filing thresholds.
     """
 
-    divr_array: ArrayDouble
-    """Diversion ratio between the merging firms"""
+    @nth_firm_share.default
+    def __nfsd(_i: MarketSampleData) -> ArrayINT:
+        e_ = np.empty_like(_i.frmshr_array[:, :1], float)
+        e_.fill(np.nan)
+        return e_
 
-    hhi_post: ArrayDouble
+    hhi_post: ArrayDouble = field(eq=cmp_using(np.array_equal))
     """Post-merger change in Herfindahl-Hirschmann Index (HHI)"""
 
-    hhi_delta: ArrayDouble
-    """Change in HHI from combination of merging firms"""
+    @hhi_post.default
+    def __hpd(_i: MarketSampleData) -> ArrayINT:
+        e_ = np.empty_like(_i.frmshr_array[:, :1], float)
+        e_.fill(np.nan)
+        return e_
 
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class ShareDataSample:
     """Container for generated market shares.
 
@@ -507,7 +513,7 @@ class ShareDataSample:
     mktshr_array: ArrayDouble
     """All-firm shares (with two merging firms)"""
 
-    fcounts: ArrayBIGINT
+    fcounts: ArrayINT
     """All-firm-count for each draw"""
 
     nth_firm_share: ArrayDouble
@@ -517,7 +523,7 @@ class ShareDataSample:
     """Converts market shares to choice probabilities by multiplication."""
 
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class PriceDataSample:
     """Container for generated price array, and related."""
 
@@ -528,7 +534,7 @@ class PriceDataSample:
     """Flags draws as meeting HSR filing thresholds or not"""
 
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class MarginDataSample:
     """Container for generated margin array and related MNL test array."""
 
@@ -549,14 +555,15 @@ class MarginDataSample:
     """
 
 
+@this_yaml.register_class
 @enum.unique
-class INVResolution(enum.StrEnum):
+class INVResolution(str, EnumYAMLized):
     CLRN = "clearance"
     ENFT = "enforcement"
     BOTH = "investigation"
 
 
-@frozen_attrs
+@frozen
 class UPPTestRegime:
     """Configuration for UPP tests."""
 
@@ -576,7 +583,7 @@ class UPPTestRegime:
     """Aggregator for diversion ratio test."""
 
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class UPPTestsRaw:
     """Container for arrays marking test failures and successes
 
@@ -601,7 +608,7 @@ class UPPTestsRaw:
     """True if IPR (partial price-simulation) estimate meets criterion"""
 
 
-@dataclass(slots=True, frozen=True)
+@frozen
 class UPPTestsCounts:
     """Counts of markets resolved as specified
 
@@ -612,9 +619,9 @@ class UPPTestsCounts:
 
     """
 
-    by_firm_count: ArrayBIGINT
-    by_delta: ArrayBIGINT
-    by_conczone: ArrayBIGINT
+    by_firm_count: ArrayBIGINT = field(eq=cmp_using(eq=np.array_equal))
+    by_delta: ArrayBIGINT = field(eq=cmp_using(eq=np.array_equal))
+    by_conczone: ArrayBIGINT = field(eq=cmp_using(eq=np.array_equal))
     """Zones are "unoncentrated", "moderately concentrated", and "highly concentrated",
     with futher detail by HHI and ΔHHI for mergers in the "unconcentrated" and
     "moderately concentrated" zones. See
@@ -624,50 +631,5 @@ class UPPTestsCounts:
     """
 
 
-# https://stackoverflow.com/questions/54668000
-class DataclassInstance(Protocol):
-    """Generic dataclass-instance"""
-
-    __dataclass_fields__: ClassVar
-
-
-for _typ in (
-    PriceSpec,
-    SHRDistribution,
-    PCMDistribution,
-    FM2Constraint,
-    SSZConstant,
-    INVResolution,
-):
-    #  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)
-            ),
-        ),
-    )
-
-for _typ in (MarketSampleData, SeedSequenceData, UPPTestsCounts):
-    _, _ = (
-        this_yaml.representer.add_representer(
-            _typ,
-            lambda _r, _d: _r.rpresent_mapping(
-                f"!{_d.__class__.__name__}",
-                {_a: getattr(_d, _a) for _a in _d.__dataclass_fields__},
-            ),
-        ),
-        this_yaml.constructor.add_constructor(
-            _typ,
-            lambda _c, _n: globals().get(
-                _n.tag.lstrip("!")(**_c.construct_mapping(_n))
-            ),
-        ),
-    )
+for _typ in (SeedSequenceData, ShareSpec, PCMSpec, UPPTestsCounts, UPPTestRegime):
+    yamelize_attrs(_typ)