fix lint and tests

Author: gianlucadetommaso

inga/scm/random.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Callable, Iterable, Literal
+from typing import Callable, Iterable, Literal, cast
 
 import random
 import torch
@@ -108,8 +108,7 @@ def to_dict(self) -> dict[str, object]:
             "temperature": self._temperature,
             "bias": list(self._bias),
             "gaussian_parent_weights": {
-                key: list(value)
-                for key, value in self._gaussian_parent_weights.items()
+                key: list(value) for key, value in self._gaussian_parent_weights.items()
             },
             "categorical_parent_weights": {
                 key: [list(row) for row in value]
@@ -120,28 +119,84 @@ def to_dict(self) -> dict[str, object]:
 
     @classmethod
     def from_dict(cls, payload: dict[str, object]) -> "RandomCategoricalVariable":
+        parent_names_raw = payload.get("parent_names", [])
+        if not isinstance(parent_names_raw, list):
+            raise ValueError("`parent_names` must be a list.")
+
+        parent_kinds_raw = payload.get("parent_kinds", {})
+        if not isinstance(parent_kinds_raw, dict):
+            raise ValueError("`parent_kinds` must be a dict.")
+        parent_kinds: dict[str, Literal["gaussian", "categorical"]] = {}
+        for key, value in parent_kinds_raw.items():
+            value_str = str(value)
+            if value_str not in ("gaussian", "categorical"):
+                raise ValueError(
+                    "`parent_kinds` values must be 'gaussian' or 'categorical'."
+                )
+            parent_kinds[str(key)] = cast(
+                Literal["gaussian", "categorical"],
+                value_str,
+            )
+
+        bias_raw = payload.get("bias", [])
+        if not isinstance(bias_raw, list):
+            raise ValueError("`bias` must be a list.")
+
+        gaussian_parent_weights_raw = payload.get("gaussian_parent_weights", {})
+        if not isinstance(gaussian_parent_weights_raw, dict):
+            raise ValueError("`gaussian_parent_weights` must be a dict.")
+        gaussian_parent_weights: dict[str, list[float]] = {}
+        for key, values in gaussian_parent_weights_raw.items():
+            if not isinstance(values, list):
+                raise ValueError(
+                    "`gaussian_parent_weights` entries must be lists of floats."
+                )
+            gaussian_parent_weights[str(key)] = [float(v) for v in values]
+
+        categorical_parent_weights_raw = payload.get("categorical_parent_weights", {})
+        if not isinstance(categorical_parent_weights_raw, dict):
+            raise ValueError("`categorical_parent_weights` must be a dict.")
+        categorical_parent_weights: dict[str, list[list[float]]] = {}
+        for key, values in categorical_parent_weights_raw.items():
+            if not isinstance(values, list):
+                raise ValueError(
+                    "`categorical_parent_weights` entries must be lists of lists."
+                )
+            rows: list[list[float]] = []
+            for row in values:
+                if not isinstance(row, list):
+                    raise ValueError(
+                        "`categorical_parent_weights` entries must be lists of lists."
+                    )
+                rows.append([float(v) for v in row])
+            categorical_parent_weights[str(key)] = rows
+
+        transforms_raw = payload.get("transforms", [])
+        if not isinstance(transforms_raw, list):
+            raise ValueError("`transforms` must be a list.")
+
+        num_categories_raw = payload.get("num_categories")
+        if num_categories_raw is None:
+            raise ValueError("`num_categories` is required.")
+        if not isinstance(num_categories_raw, (int, float, str)):
+            raise ValueError("`num_categories` must be int-like.")
+
+        temperature_raw = payload.get("temperature")
+        if temperature_raw is None:
+            raise ValueError("`temperature` is required.")
+        if not isinstance(temperature_raw, (int, float, str)):
+            raise ValueError("`temperature` must be float-like.")
+
         return cls(
             name=str(payload["name"]),
-            parent_names=[str(name) for name in payload.get("parent_names", [])],
-            parent_kinds={
-                str(key): str(value)  # type: ignore[dict-item]
-                for key, value in dict(payload.get("parent_kinds", {})).items()
-            },
-            num_categories=int(payload["num_categories"]),
-            temperature=float(payload["temperature"]),
-            bias=[float(value) for value in list(payload["bias"])],
-            gaussian_parent_weights={
-                str(key): [float(v) for v in list(values)]
-                for key, values in dict(payload.get("gaussian_parent_weights", {})).items()
-            },
-            categorical_parent_weights={
-                str(key): [
-                    [float(v) for v in list(row)]
-                    for row in list(values)
-                ]
-                for key, values in dict(payload.get("categorical_parent_weights", {})).items()
-            },
-            transforms=[str(name) for name in list(payload.get("transforms", []))],
+            parent_names=[str(name) for name in parent_names_raw],
+            parent_kinds=parent_kinds,
+            num_categories=int(num_categories_raw),
+            temperature=float(temperature_raw),
+            bias=[float(value) for value in bias_raw],
+            gaussian_parent_weights=gaussian_parent_weights,
+            categorical_parent_weights=categorical_parent_weights,
+            transforms=[str(name) for name in transforms_raw],
         )
 
 
@@ -232,6 +287,50 @@ def f_mean(parents: dict[str, Tensor]) -> Tensor:
     return f_mean
 
 
+def _build_mixed_parent_mean(
+    parent_names: list[str],
+    intercept: float,
+    scalar_coefs: dict[str, float],
+    parent_projections: dict[str, list[float]],
+    transforms: list[Callable[[Tensor], Tensor]] | None,
+) -> Callable[[dict[str, Tensor]], Tensor]:
+    """Build Gaussian mean function supporting both scalar and categorical parents."""
+    parent_names_local = list(parent_names)
+    intercept_local = intercept
+    scalar_coefs_local = dict(scalar_coefs)
+    parent_projections_local = {
+        name: list(weights) for name, weights in parent_projections.items()
+    }
+    transforms_local = transforms
+
+    def _as_scalar(parent_name: str, parent_value: Tensor) -> Tensor:
+        if parent_value.ndim <= 1:
+            return parent_value
+        projection = torch.tensor(
+            parent_projections_local[parent_name],
+            device=parent_value.device,
+            dtype=parent_value.dtype,
+        )
+        return (parent_value * projection).sum(dim=-1)
+
+    def mixed_parent_mean(base_parents: dict[str, Tensor]) -> Tensor:
+        if not parent_names_local or not base_parents:
+            base = torch.tensor(intercept_local)
+        else:
+            ref = next(iter(base_parents.values()))
+            base = (
+                torch.zeros_like(ref[..., 0] if ref.ndim > 1 else ref) + intercept_local
+            )
+            for parent_name in parent_names_local:
+                parent_scalar = _as_scalar(parent_name, base_parents[parent_name])
+                base = base + scalar_coefs_local[parent_name] * parent_scalar
+        if transforms_local:
+            return _compose_transforms(base, transforms_local)
+        return base
+
+    return mixed_parent_mean
+
+
 def _sample_parents(
     rng: random.Random,
     previous_names: list[str],
@@ -320,10 +419,7 @@ def random_scm(config: RandomSCMConfig) -> SCM:
                 rng.uniform(*config.categorical_bias_range)
                 for _ in range(num_categories)
             ]
-            parent_kinds = {
-                parent: variable_kinds[parent]
-                for parent in parents
-            }
+            parent_kinds = {parent: variable_kinds[parent] for parent in parents}
 
             gaussian_parent_weights: dict[str, list[float]] = {}
             categorical_parent_weights: dict[str, list[list[float]]] = {}
@@ -343,8 +439,10 @@ def random_scm(config: RandomSCMConfig) -> SCM:
                         for _ in range(parent_categories)
                     ]
 
-            transform_names = (
-                _sample_transform_names(rng) if rng.random() < config.nonlinear_prob else []
+            categorical_transform_names = (
+                _sample_transform_names(rng)
+                if rng.random() < config.nonlinear_prob
+                else []
             )
 
             variables.append(
@@ -357,76 +455,72 @@ def random_scm(config: RandomSCMConfig) -> SCM:
                     bias=categorical_bias,
                     gaussian_parent_weights=gaussian_parent_weights,
                     categorical_parent_weights=categorical_parent_weights,
-                    transforms=transform_names,
+                    transforms=categorical_transform_names,
                 )
             )
             variable_kinds[name] = "categorical"
             variable_categories[name] = num_categories
             continue
 
         use_nonlinear = rng.random() < config.nonlinear_prob
-        transform_names = _sample_transform_names(rng) if use_nonlinear else None
-        transforms = resolve_transforms(transform_names) if transform_names else None
+        gaussian_transform_names: list[str] | None = (
+            _sample_transform_names(rng) if use_nonlinear else None
+        )
+        transforms = (
+            resolve_transforms(gaussian_transform_names)
+            if gaussian_transform_names
+            else None
+        )
+
+        parent_names_local = list(parents)
+        intercept_local = intercept
 
         parent_projections: dict[str, list[float]] = {
             parent: [
                 rng.uniform(*config.coef_range)
                 for _ in range(variable_categories[parent])
             ]
-            for parent in parents
+            for parent in parent_names_local
             if variable_kinds.get(parent) == "categorical"
         }
 
-        def _as_scalar(parent_name: str, parent_value: Tensor) -> Tensor:
-            if parent_value.ndim <= 1:
-                return parent_value
-            projection = torch.tensor(
-                parent_projections[parent_name],
-                device=parent_value.device,
-                dtype=parent_value.dtype,
-            )
-            return (parent_value * projection).sum(dim=-1)
-
-        scalar_coefs = {parent: rng.uniform(*config.coef_range) for parent in parents}
-
-        def mixed_parent_mean(base_parents: dict[str, Tensor]) -> Tensor:
-            if not parents:
-                base = torch.tensor(intercept)
-            else:
-                ref = next(iter(base_parents.values()))
-                base = torch.zeros_like(ref[..., 0] if ref.ndim > 1 else ref) + intercept
-                for parent_name in parents:
-                    parent_scalar = _as_scalar(parent_name, base_parents[parent_name])
-                    base = base + scalar_coefs[parent_name] * parent_scalar
-            if transforms:
-                return _compose_transforms(base, transforms)
-            return base
+        scalar_coefs = {
+            parent: rng.uniform(*config.coef_range) for parent in parent_names_local
+        }
+        mixed_parent_mean = _build_mixed_parent_mean(
+            parent_names=parent_names_local,
+            intercept=intercept_local,
+            scalar_coefs=scalar_coefs,
+            parent_projections=parent_projections,
+            transforms=transforms,
+        )
+        gaussian_transform_names_local = gaussian_transform_names
 
         has_categorical_parent = any(
             variable_kinds.get(parent_name) == "categorical"
-            for parent_name in parents
+            for parent_name in parent_names_local
         )
 
         if use_nonlinear or has_categorical_parent:
             variables.append(
                 FunctionalVariable(
                     name=name,
-                    parent_names=parents,
+                    parent_names=parent_names_local,
                     sigma=sigma,
                     f_mean=mixed_parent_mean,
                     coefs=scalar_coefs,
-                    intercept=intercept,
-                    transforms=transform_names,
+                    intercept=intercept_local,
+                    transforms=gaussian_transform_names_local,
                 )
             )
         else:
             variables.append(
                 LinearVariable(
                     name=name,
-                    parent_names=parents,
+                    parent_names=parent_names_local,
                     sigma=sigma,
                     coefs=scalar_coefs,
-                    intercept=intercept,
+                    intercept=intercept_local,
                 )
             )
         variable_kinds[name] = "gaussian"