add simple sub-state model accessor and evaluation

2 files changed, 71 insertions, 2 deletions

diff --git a/bin/analyze-archive.py b/bin/analyze-archive.py
index 1b37b89..40f11e8 100755
--- a/bin/analyze-archive.py
+++ b/bin/analyze-archive.py
@@ -815,6 +815,7 @@ if __name__ == "__main__":
     if len(show_models):
         print("--- param model ---")
 
+    # get_fitted_sub -> with sub-state detection and modeling
     param_model, param_info = model.get_fitted(
         safe_functions_enabled=safe_functions_enabled
     )
@@ -903,6 +904,30 @@ if __name__ == "__main__":
                             param_info(trans, attribute)["function"].model_args,
                         )
                     )
+        if args.with_substates:
+            for submodel in model.submodel_by_nc.values():
+                sub_param_model, sub_param_info = submodel.get_fitted()
+                for substate in submodel.states():
+                    for subattribute in submodel.attributes(substate):
+                        if sub_param_info(substate, subattribute):
+                            print(
+                                "{:10s} {:15s}: {}".format(
+                                    substate,
+                                    subattribute,
+                                    sub_param_info(substate, subattribute)[
+                                        "function"
+                                    ].model_function,
+                                )
+                            )
+                            print(
+                                "{:10s} {:15s}  {}".format(
+                                    "",
+                                    "",
+                                    sub_param_info(substate, subattribute)[
+                                        "function"
+                                    ].model_args,
+                                )
+                            )
 
     if xv_method == "montecarlo":
         analytic_quality = xv.montecarlo(lambda m: m.get_fitted()[0], xv_count)
@@ -940,6 +965,18 @@ if __name__ == "__main__":
         model_quality_table(
             [static_quality, analytic_quality, lut_quality], [None, param_info, None]
         )
+        if args.with_substates:
+            for submodel in model.submodel_by_nc.values():
+                sub_static_model = submodel.get_static()
+                sub_static_quality = submodel.assess(sub_static_model)
+                sub_lut_model = submodel.get_param_lut()
+                sub_lut_quality = submodel.assess(sub_lut_model)
+                sub_param_model, sub_param_info = submodel.get_fitted()
+                sub_analytic_quality = submodel.assess(sub_param_model)
+                model_quality_table(
+                    [sub_static_quality, sub_analytic_quality, sub_lut_quality],
+                    [None, sub_param_info, None],
+                )
 
     if "overall" in show_quality or "all" in show_quality:
         print("overall state static/param/lut MAE assuming equal state distribution:")
diff --git a/lib/model.py b/lib/model.py
index 9ac4560..44272c8 100644
--- a/lib/model.py
+++ b/lib/model.py
@@ -806,7 +806,6 @@ class PTAModel(AnalyticModel):
 
             self.pelt = PELT(**pelt)
             self.find_substates()
-            print(self.submodel_by_nc)
         else:
             self.pelt = None
 
@@ -828,6 +827,37 @@ class PTAModel(AnalyticModel):
             for key in elem["attributes"]:
                 elem[key] = np.array(elem[key])
 
+    def get_fitted_sub(self, use_mean=False, safe_functions_enabled=False):
+
+        param_model_getter, param_info_getter = self.get_fitted(
+            use_mean=use_mean, safe_functions_enabled=safe_functions_enabled
+        )
+
+        def model_getter(name, key, **kwargs):
+            if key != "power":
+                return param_model_getter(name, key, **kwargs)
+
+            try:
+                substate_count = round(param_model_getter(name, "substate_count"))
+            except KeyError:
+                return param_model_getter(name, key, **kwargs)
+            if substate_count == 1:
+                return param_model_getter(name, key, **kwargs)
+
+            cumulative_energy = 0
+            total_duration = 0
+            substate_model, _ = self.submodel_by_nc[(name, substate_count)].get_fitted()
+            for i in range(substate_count):
+                sub_name = f"{name}.{i+1}({substate_count})"
+                cumulative_energy += substate_model(
+                    sub_name, "duration", **kwargs
+                ) * substate_model(sub_name, "power", **kwargs)
+                total_duration += substate_model(sub_name, "duration", **kwargs)
+
+            return cumulative_energy / total_duration
+
+        return model_getter, param_info_getter
+
     # This heuristic is very similar to the "function is not much better than
     # median" checks in get_fitted. So far, doing it here as well is mostly
     # a performance and not an algorithm quality decision.
@@ -991,7 +1021,9 @@ class PTAModel(AnalyticModel):
                 # data units are s / W, models use µs / µW
                 durations.extend(np.array(run[substate_index]["duration"]) * 1e6)
                 powers.extend(np.array(run[substate_index]["power"]) * 1e6)
-                param_values.extend([param for i in run[substate_index]["duration"]])
+                param_values.extend(
+                    [list(param) for i in run[substate_index]["duration"]]
+                )
 
             by_name[sub_name] = {
                 "isa": "state",