Get rid of model_dict.

Author: hmgaudecker

docs/how_to_guides/model_specs.md

@@ -1,7 +1,7 @@
 # Model Specifications
 
-Models can be specified using Python dataclasses or dictionaries. The dataclass approach
-is recommended for type safety and IDE support.
+Models are specified using Python dataclasses. YAML files can be loaded and converted to
+`ModelSpec` via `from_dict()`.
 
 ## Using Dataclasses (Recommended)
 
@@ -41,19 +41,22 @@ model = ModelSpec(
 )
 ```
 
-## Using Dictionaries
+## Loading from YAML
 
-For backwards compatibility and interoperability with YAML/JSON files, models can also
-be specified as dictionaries:
+Models can be loaded from YAML files by parsing the YAML and converting to a `ModelSpec`
+via `from_dict()`:
 
 ```python
 import yaml
+from skillmodels import ModelSpec
 
 with open("model.yaml") as f:
-    model = yaml.safe_load(f)
+    raw = yaml.safe_load(f)
+
+model = ModelSpec.from_dict(raw)
 ```
 
-The dictionary structure mirrors the dataclass structure:
+The YAML dictionary structure mirrors the dataclass structure:
 
 ```yaml
 factors:

pixi.lock

@@ -20,7 +20,6 @@ classifiers = [
 ]
 dependencies = [
   "dags",
-  "frozendict",
   "jax>=0.8",
   "numpy",
   "pandas",

pyproject.toml

@@ -1,5 +1,7 @@
 """Functions to validate model specifications."""
 
+from collections.abc import Mapping
+
 import numpy as np
 
 from skillmodels.types import Anchoring, Dimensions, Labels
@@ -92,8 +94,8 @@ def _check_anchoring(anchoring: Anchoring) -> list[str]:
     report = []
     if not isinstance(anchoring.anchoring, bool):
         report.append("anchoring.anchoring must be a bool.")
-    if not isinstance(anchoring.outcomes, dict):
-        report.append("anchoring.outcomes must be a dict")
+    if not isinstance(anchoring.outcomes, Mapping):
+        report.append("anchoring.outcomes must be a Mapping")
     else:
         variables = list(anchoring.outcomes.values())
         for var in variables:

src/skillmodels/check_model.py

@@ -389,7 +389,7 @@ def _get_anchoring_constraints(
         ind_tups = []
         for period in periods:
             for factor in anchoring_info.factors:
-                outcome = anchoring_info.outcomes[factor]  # ty: ignore[invalid-argument-type]
+                outcome = anchoring_info.outcomes[factor]
                 meas = f"{outcome}_{factor}"
                 ind_tups.append(("loadings", period, meas, factor))
 
@@ -431,7 +431,7 @@ def _get_constraints_for_augmented_periods(
         # look counterintuitive...
         aug_period_meas_type_to_constrain = (
             MeasurementType.STATES
-            if endogenous_factors_info.factor_info[factor].is_state  # ty: ignore[invalid-argument-type]
+            if endogenous_factors_info.factor_info[factor].is_state
             else MeasurementType.ENDOGENOUS_FACTORS
         )
         aug_period_meas_types = (

src/skillmodels/constraints.py

@@ -134,7 +134,7 @@ def plot_correlation_heatmap(
 
 def get_measurements_corr(
     data: pd.DataFrame,
-    model: dict | ModelSpec,
+    model: ModelSpec,
     factors: list[str] | tuple[str, ...] | str | None,
     periods: float | list[int] | None,
 ) -> pd.DataFrame:
@@ -178,7 +178,7 @@ def get_measurements_corr(
 
 def get_quasi_scores_corr(
     data: pd.DataFrame,
-    model: dict | ModelSpec,
+    model: ModelSpec,
     factors: list[str] | tuple[str, ...] | str | None,
     periods: float | list[int] | None,
 ) -> pd.DataFrame:
@@ -226,7 +226,7 @@ def get_quasi_scores_corr(
 def get_scores_corr(
     data: pd.DataFrame,
     params: pd.DataFrame,
-    model: dict | ModelSpec,
+    model: ModelSpec,
     factors: list[str] | tuple[str, ...] | str | None,
     periods: float | list[int] | None,
 ) -> pd.DataFrame:

src/skillmodels/correlation_heatmap.py

@@ -15,7 +15,7 @@
 
 
 def get_filtered_states(
-    model: dict | ModelSpec,
+    model: ModelSpec,
     data: pd.DataFrame,
     params: pd.DataFrame,
 ) -> dict[str, dict[str, Any]]:
@@ -54,7 +54,7 @@ def get_filtered_states(
 
 def anchor_states_df(
     states_df: pd.DataFrame,
-    model: dict | ModelSpec,
+    model: ModelSpec,
     params: pd.DataFrame,
     *,
     use_aug_period: bool,

src/skillmodels/filtered_states.py

@@ -32,14 +32,14 @@
 
 
 def get_maximization_inputs(
-    model: dict | ModelSpec,
+    model: ModelSpec,
     data: pd.DataFrame,
     split_dataset: int = 1,
 ) -> dict[str, Any]:
     """Create inputs for optimagic's maximize function.
 
     Args:
-        model: The model specification, either as a dict or ModelSpec instance.
+        model: The model specification as a ModelSpec instance.
             See: :ref:`model_specs`
         data: dataset in long format.
         split_dataset(Int): Controls into how many sclices to split the dataset

src/skillmodels/maximization_inputs.py

@@ -2,16 +2,15 @@
 
 This module provides frozen dataclasses for defining model specifications
 in a type-safe, immutable manner. All collections use immutable types
-(tuples, frozendict) to ensure the specification cannot be accidentally modified.
+(tuples, MappingProxyType) to ensure the specification cannot be accidentally
+modified.
 """
 
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from types import MappingProxyType
 from typing import Self
 
-from frozendict import frozendict
-
 
 @dataclass(frozen=True)
 class Normalizations:
@@ -25,15 +24,15 @@ class Normalizations:
 
     """
 
-    loadings: tuple[frozendict[str, float], ...]
-    intercepts: tuple[frozendict[str, float], ...]
+    loadings: tuple[MappingProxyType[str, float], ...]
+    intercepts: tuple[MappingProxyType[str, float], ...]
 
     @classmethod
     def from_dict(cls, d: dict) -> Self:
         """Create Normalizations from a dictionary specification."""
         return cls(
-            loadings=tuple(frozendict(x) for x in d["loadings"]),
-            intercepts=tuple(frozendict(x) for x in d["intercepts"]),
+            loadings=tuple(MappingProxyType(x) for x in d["loadings"]),
+            intercepts=tuple(MappingProxyType(x) for x in d["intercepts"]),
         )
 
     def to_dict(self) -> dict:
@@ -169,8 +168,8 @@ def to_dict(self) -> dict:
         return result
 
 
-def _default_empty_frozendict() -> frozendict[str, str]:
-    return frozendict({})
+def _default_empty_mapping_proxy() -> MappingProxyType[str, str]:
+    return MappingProxyType({})
 
 
 @dataclass(frozen=True)
@@ -186,7 +185,9 @@ class AnchoringSpec:
 
     """
 
-    outcomes: frozendict[str, str] = field(default_factory=_default_empty_frozendict)
+    outcomes: MappingProxyType[str, str] = field(
+        default_factory=_default_empty_mapping_proxy,
+    )
     free_controls: bool = False
     free_constant: bool = False
     free_loadings: bool = False
@@ -198,7 +199,7 @@ def from_dict(cls, d: dict) -> Self:
         outcomes = d.get("outcomes", {})
         ignore_constant = d.get("ignore_constant_when_anchoring", False)
         return cls(
-            outcomes=frozendict(outcomes),
+            outcomes=MappingProxyType(outcomes),
             free_controls=d.get("free_controls", False),
             free_constant=d.get("free_constant", False),
             free_loadings=d.get("free_loadings", False),

src/skillmodels/model_spec.py

@@ -177,7 +177,7 @@ def _add_copies_of_anchoring_outcome(
 ) -> pd.DataFrame:
     df = df.copy()
     for factor in anchoring_info.factors:
-        outcome = anchoring_info.outcomes[factor]  # ty: ignore[invalid-argument-type]
+        outcome = anchoring_info.outcomes[factor]
         df[f"{outcome}_{factor}"] = df[outcome]
     return df