add ParamType helper class

2 files changed, 48 insertions, 13 deletions

diff --git a/lib/model.py b/lib/model.py
index f01ae5d..f25e44f 100644
--- a/lib/model.py
+++ b/lib/model.py
@@ -6,6 +6,7 @@ import os
 from .automata import PTA, ModelAttribute
 from .functions import StaticFunction, SubstateFunction, SplitFunction
 from .parameters import (
+    ParamType,
     ParallelParamStats,
     ParamStats,
     codependent_param_dict,
@@ -127,10 +128,14 @@ class AnalyticModel:
             self._num_args = _num_args_from_by_name(by_name)
 
         self.distinct_param_values_by_name = dict()
+        self.param_type_by_name = dict()
         for name in self.names:
             self.distinct_param_values_by_name[name] = distinct_param_values(
                 by_name[name]["param"]
             )
+            self.param_type_by_name[name] = ParamType(
+                self.distinct_param_values_by_name[name], values_are_distinct=True
+            )
 
         self.fit_done = False
 
@@ -439,24 +444,20 @@ class AnalyticModel:
                 "scalar": 0,
                 "enum": 0,
             }
-            for param_index, param_values in enumerate(
-                self.distinct_param_values_by_name[name]
-            ):
-                if None in param_values:
-                    none_adj = -1
-                else:
-                    none_adj = 0
-                value_count = len(param_values) + none_adj
-                if value_count == 0:
+            for param_index in range(len(self.parameters)):
+                param_type = self.param_type_by_name[name][param_index]
+                if param_type == ParamType.UNSET:
                     param_data["unset"] += 1
-                elif value_count == 1:
+                elif param_type == ParamType.USELESS:
                     param_data["useless"] += 1
-                elif value_count == 2:
+                elif param_type == ParamType.BOOLEAN:
                     param_data["boolean"] += 1
-                elif all(map(lambda x: x is None or is_numeric(x), param_values)):
+                elif param_type == ParamType.SCALAR:
                     param_data["scalar"] += 1
-                else:
+                elif param_type == ParamType.ENUM:
                     param_data["enum"] += 1
+                else:
+                    raise RuntimeError(f"Unknown param_type: {param_type}")
             ret[f"paramcount/{name}/useful"] = (
                 param_data["boolean"] + param_data["scalar"] + param_data["enum"]
             )
@@ -663,10 +664,14 @@ class PTAModel(AnalyticModel):
         self.ignore_trace_indexes = ignore_trace_indexes
 
         self.distinct_param_values_by_name = dict()
+        self.param_type_by_name = dict()
         for name in self.names:
             self.distinct_param_values_by_name[name] = distinct_param_values(
                 by_name[name]["param"]
             )
+            self.param_type_by_name[name] = ParamType(
+                self.distinct_param_values_by_name[name], values_are_distinct=True
+            )
 
         self.fit_done = False
 
diff --git a/lib/parameters.py b/lib/parameters.py
index 40d0ba5..1b46a8e 100644
--- a/lib/parameters.py
+++ b/lib/parameters.py
@@ -308,6 +308,36 @@ def _all_params_are_numeric(data, param_idx):
     return all(map(is_numeric, param_values))
 
 
+class ParamType(dict):
+    UNSET = 0
+    USELESS = 1
+    BOOLEAN = 2
+    SCALAR = 3
+    ENUM = 4
+
+    def __init__(self, param_values, values_are_distinct=False):
+        if values_are_distinct:
+            distinct_values = param_values
+        else:
+            distinct_values = distinct_param_values(param_values)
+        for param_index, param_values in enumerate(distinct_values):
+            if None in param_values:
+                none_adj = -1
+            else:
+                none_adj = 0
+            value_count = len(param_values) + none_adj
+            if value_count == 0:
+                self[param_index] = self.UNSET
+            elif value_count == 1:
+                self[param_index] = self.USELESS
+            elif value_count == 2:
+                self[param_index] = self.BOOLEAN
+            elif all(map(lambda n: n is None or is_numeric(n), param_values)):
+                self[param_index] = self.SCALAR
+            else:
+                self[param_index] = self.ENUM
+
+
 class ParallelParamStats:
     def __init__(self):
         self.queue = list()