Refactor stan builder, add loglik inside stan, share step_size scheme

Author: hyunjimoon

metric_checkpoint_round_0_chain_1_pridraws_0.json

@@ -0,0 +1 @@
+{"inv_metric": [0.00890483, 0.0146411, 0.0180988, 0.0173486, 0.0162696, 0.0185366, 0.0186761, 0.0176091, 0.0177719, 0.0194737, 0.0269911, 0.0241012, 0.0270683, 0.0305894, 0.0295884, 0.0285179, 0.0270491, 0.0247198, 0.00094971]}

stanify/builders/stan_block_builder.py

@@ -482,24 +482,18 @@ def build_block(self, hier_est_param_names):
         if self.precision_context.R == 1:
             for statement in self.stan_model_context.sample_statements:
                 if statement.distribution_type != statement.assignment_dist:
-                    code += f"{statement.lhs_expr} ~ {statement.distribution_type}({', '.join([str(arg) for arg in statement.distribution_args])});\n"
+                    code += f"{adj_expr(statement)};\n"
         else:
             for statement in self.stan_model_context.sample_statements:
-                param_name = statement.lhs_expr
-                if param_name in hier_est_param_names:
+                if statement.lhs_expr in hier_est_param_names:
                     dist_code = "rep_vector(" + f'{statement.distribution_args[0]}, R), ' + f"{', '.join(statement.distribution_args[1:])}"
-                    code += f"{param_name} ~ {statement.distribution_type}({dist_code});\n"
-
-                elif param_name in self.stan_model_context.obs_integ_outcome_vector_names:
-                    code += "for (r in 1:R)\n"
-                    code.indent_level += 1
-                    dist_code = f'{param_name}'[:-4] + "[:, r], " + f"{', '.join(statement.distribution_args[1:])}"
-                    code += f"{param_name}[:, r] ~ {statement.distribution_type}({dist_code});\n"
-                    code.indent_level -= 1
+                    code += f"{statement.lhs_expr} ~ {statement.distribution_type}({dist_code});\n"
 
+                elif statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
+                    code += f"{adj_expr(statement, is_hier=True)};\n"
                 else:
                     if statement.distribution_type != statement.assignment_dist:
-                        code += f"{param_name} ~ {statement.distribution_type}({', '.join([str(arg) for arg in statement.distribution_args])});\n"
+                        code += f"{adj_expr(statement)};\n"
         code.indent_level -= 1
         code += "}\n"
         #TODO @Dashadower what is the diff btw classes that has its own code vs not (self.code VS return str(code))
@@ -511,25 +505,25 @@ def __init__(self, precision_context: "PrecisionContext", stan_model_context: "S
         self.precision_context = precision_context
         self.stan_model_context = stan_model_context
         self.vensim_model_context = vensim_model_context
-        # TODO @Dashadower how to write message, if some target_simulated_vector_names is not in vensim_integ_outcome (inconsistency btw user-defined and vensim syntax)
+        # TODO @Dashadower how to write message, if some target_sim_vector_names is not in vensim_integ_outcome (inconsistency btw user-defined and vensim syntax)
         integ_outcome_vector_names = set(self.stan_model_context.target_integ_outcome_vector_names) & set(self.vensim_model_context.integ_outcome_vector_names)
 
     def build_block(self, hier_est_param_names, transformed_parameters_code: str = ""):
         self.code = IndentedString()
         self.code += "generated quantities{\n"
         self.code.indent_level += 1
-        self.build_param_rng_functions(hier_est_param_names)
+        self.build_param_pri_pred_functions(hier_est_param_names)
         self.code += "\n"
         self.code.add_raw(transformed_parameters_code, ignore_indent=True)
         self.code += "\n"
-        self.build_obs_rng_functions()
+        self.build_data_pri_pred_functions()
         self.code.indent_level -= 1
         self.code += "}\n"
 
         return str(self.code)
 
 
-    def build_param_rng_functions(self, hier_est_param_names):
+    def build_param_pri_pred_functions(self, hier_est_param_names):
 
         ignored_variables = set(self.stan_model_context.stan_data.keys()).union(
             set(self.vensim_model_context.integ_outcome_vector_names))
@@ -557,40 +551,29 @@ def build_param_rng_functions(self, hier_est_param_names):
                     if statement.init_state:
                         param_name = param_name + "__init"
                     if param_name in hier_est_param_names:
-
-
                         dist_code = "rep_vector(" + f'{statement.distribution_args[0]}, R), ' + f"{', '.join(statement.distribution_args[1:])}"
                         self.code += f"real {param_name}[R] =  {statement.distribution_type}_rng({dist_code});\n"
                     else:
-                        self.code += f"real {param_name} = {statement.distribution_type}_rng({', '.join(statement.distribution_args)});\n"
+                        self.code += f"real {adj_expr(statement, is_pri_pred=True)};\n"
                     processed_statements.add(statement)
 
-    def build_obs_rng_functions(self):
+    def build_data_pri_pred_functions(self):
         if self.precision_context.R == 1:
             self.code += "// Define and assign generated value to observed vector (matching vector)\n"
-
             for statement in self.stan_model_context.sample_statements:
                 if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
-                    vec_name = statement.lhs_expr
-                    self.code += f"array [N] real {vec_name} = {statement.distribution_type}_rng({', '.join(statement.distribution_args)});\n"
+                    self.code += f"array [N] real {adj_expr(statement, is_pri_pred=True)};\n"
         else:
             self.code += "// Define observed vector (matching vector)\n"
             for statement in self.stan_model_context.sample_statements:
                 if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
                     self.code += f"array[N] vector[R] {statement.lhs_expr};\n"
 
             self.code += "// Assign generated value to observed vector (matching vector)\n"
-            self.code += "for (r in 1:R){\n"
-            self.code.indent_level += 1
             for statement in self.stan_model_context.sample_statements:
-                #TODO @Dashadower statement.lhs_variable vs .lhs_expr
                 if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
-                    vec_name = statement.lhs_expr
-                    dist_code = f'{vec_name}'[:-4] + "[:, r], " + f"{', '.join(statement.distribution_args[1:])}"
-                    self.code += f"{vec_name}[:, r] = {statement.distribution_type}_rng({dist_code});\n"
+                    self.code += f"{adj_expr(statement, is_pri_pred=True, is_hier=True)};\n"
             # link(alpha) ~ N(0,1); link(alpha) is expr, alpha is var
-            self.code.indent_level -= 1
-            self.code += "}\n"
 
 class Data2DrawsStanGQBuilder():
     def __init__(self, precision_context: "PrecisionContext", stan_model_context: "StanModelContext", vensim_model_context: "VensimModelContext"):
@@ -599,75 +582,58 @@ def __init__(self, precision_context: "PrecisionContext", stan_model_context: "S
         self.vensim_model_context = vensim_model_context
 
 
-    def build_block(self):
+    def build_block(self, hier_est_param_names):
         self.code = IndentedString()
         self.code += "generated quantities{\n"
         self.code.indent_level += 1
         self.build_post_pred_rng_functions()
         self.code += "\n"
-        self.build_loglik_functions()
+        self.build_loglik_functions(hier_est_param_names)
         self.code.indent_level -= 1
         self.code += "}\n"
 
         return str(self.code)
 
-    def build_loglik_functions(self):
+    def build_loglik_functions(self, hier_est_param_names):
         self.code += "real loglik;\n"
+        self.code += "real loglik_prior;\n"
+        for tn in self.stan_model_context.target_integ_outcome_vector_names:
+            self.code += f"real loglik_{tn};\n"
 
-        if self.precision_context.R == 1:
-            for statement in self.stan_model_context.sample_statements:
-                if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
-                    param_name = statement.lhs_expr
-                    loc = statement.distribution_args[0]
-                    scale = statement.distribution_args[1]
-                    if statement.distribution_type in ["normal", "lognormal"]:
-                        self.code += f"loglik += {statement.distribution_type}_lpdf({param_name}|{loc}, {scale});\n"
-                    elif statement.distribution_type in ["neg_binom_2"]:
-                        self.code += f"loglik += {statement.distribution_type}_lpmf({param_name}|{loc}, {scale});\n"
-        else:
-            self.code += "for (r in 1:R){\n"
-            self.code.indent_level += 1
-            for statement in self.stan_model_context.sample_statements:
-                if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
-                    obs_vec_name = statement.lhs_expr
-                    target_vec_name = obs_vec_name[:-4]
-                    scale = statement.distribution_args[1]
-                    if statement.distribution_type in ["normal", "lognormal"]:
-                        self.code += f"loglik += {statement.distribution_type}_lpdf({obs_vec_name}[:, r]|{target_vec_name}[:, r], {scale});\n"
-                    elif statement.distribution_type in ["neg_binom_2"]:
-                        self.code += f"loglik += {statement.distribution_type}_lpmf({obs_vec_name}[:, r]|{target_vec_name}[:, r], {scale});\n"
-            self.code.indent_level -= 1
-            self.code += "}\n"
+        # add loglik for "matter ~ form"
+        # self.stan_model_context.all_stan_variables include "matter"
+        # 1. parameter draw ~ prior distribution
+        # X. target_sim = parameter draw
+        # 2. observed data ~ likelihood distribution (target_sim)
+
+        #is_lp_pq=True: component-wise (P for estimated parameter, Q for target_simulated)
+        for statement in self.stan_model_context.sample_statements:
+            if statement.lhs_expr in list(self.stan_model_context.all_stan_variables): #['adj_frac1[R]', 'adj_frac2', 'm_noise_scale', 'stocked_pping_obs[R]', 'stocked_ping_obs[R]']
+                # lp contribution from prior function and realized estimated parameter value
+                if statement.lhs_expr in hier_est_param_names:
+                    self.code += f"{adj_expr(statement, is_lp_tot=True, is_hier=True)};\n" #
+                    self.code += f"{adj_expr(statement, is_lp_prior=True, is_hier=True)};\n"
+
+                # lp contribution from likelihood function and observed data
+                elif statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
+                    self.code += f"{adj_expr(statement, is_lp_tot=True, is_hier=(self.precision_context.R > 1))};\n"
+                    self.code += f"{adj_expr(statement, is_lp_q=True, is_hier=(self.precision_context.R > 1))};\n"
 
-    def build_post_pred_rng_functions(self):
 
+    def build_post_pred_rng_functions(self):
         if self.precision_context.R == 1:
             self.code += "// Define and assign generated value to posterior predictive vector\n"
             for statement in self.stan_model_context.sample_statements:
                 if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
-                    # TODO @Dashadower how to use the following in the future?
-                    #  stan_type = self.stan_model_context.stan_data[statement.lhs_expr].stan_type #     stan_type = self.stan_model_context.stan_data[statement.lhs_expr].stan_type KeyError: 'prey_obs'
-                    scale = statement.distribution_args[1]
-                    # if stan_type.startswith("vector"):
-                    #     self.code += f"{stan_type} {statement.lhs_expr}_post = to_vector({statement.distribution_type}_rng({', '.join(statement.distribution_args)}));\n"
-                    # else:
-                    #     self.code += f"{stan_type} {statement.lhs_expr}_post = {statement.distribution_type}_rng({', '.join(statement.distribution_args)});\n"
-                    self.code += f"array[N] real {statement.lhs_expr}_post = {statement.distribution_type}_rng({', '.join(statement.distribution_args)});\n"
-        else:
+                    self.code += f"array[N] real {adj_expr(statement, is_post_pred=True)};\n"
+
+        elif self.precision_context.R > 1:
             self.code += "// Define observed vector (matching vector)\n"
             for statement in self.stan_model_context.sample_statements:
                 if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
                     self.code += f"array[N] vector[R] {statement.lhs_expr}_post;\n"
 
             self.code += "// Assign generated value to observed vector (matching vector)\n"
-            self.code += "for (r in 1:R){\n"
-            self.code.indent_level += 1
             for statement in self.stan_model_context.sample_statements:
-                #TODO @Dashadower statement.lhs_variable vs .lhs_expr, obs_integ_outcome_vector_names is list VS target_// is tuple so the latter doesn't work
                 if statement.lhs_expr in self.stan_model_context.obs_integ_outcome_vector_names:
-                    scale = statement.distribution_args[1]
-                   #TODO {', '.join(statement.distribution_args)} is better; posterior predictive remove _obs?? (w.o. loc, scale)
-                    self.code += f"{statement.lhs_expr}_post[:, r] = {statement.distribution_type}_rng({statement.lhs_expr[:-4]}[:, r], {scale});\n"
-            self.code.indent_level -= 1
-            self.code += "}\n"
-
+                    self.code += f"{adj_expr(statement, is_post_pred=True, is_hier=True)};\n"

stanify/builders/stan_model.py

@@ -181,24 +181,25 @@ def set_prior(self, variable_name: str, distribution_type: str, *args, lower=flo
                         if name in self.vensim_model_context.variable_names and name not in self.vensim_model_context.integ_outcome_vector_names:
                             self.stan_model_context.exposed_parameters.update(used_variable_names)
 
-            if variable_name in self.vensim_model_context.variable_names and variable_name not in self.vensim_model_context.integ_outcome_vector_names:
+            if (variable_name in self.vensim_model_context.variable_names) and (variable_name not in self.vensim_model_context.integ_outcome_vector_names):
+                #  adj_frac1_loc is excluded as it is not defined in vensim
                 self.stan_model_context.exposed_parameters.add(variable_name)
 
 
             self.stan_model_context.sample_statements.append(SamplingStatement(variable_name, distribution_type, *args, lower=lower, upper=upper, init_state=init_state))
 
 
-    def set_type(self, est_param_names: list, hier_est_param_names: list, target_simulated_vector_names: list, driving_vector_names: list, model_name: str):
+    def set_type(self, est_param_names: list, hier_est_param_names: list, target_sim_vector_names: list, driving_vector_names: list, model_name: str):
         self.est_param_names = est_param_names
         self.hier_est_param_names = hier_est_param_names
-        # TODO @Dashadower how to make target_simulated_vector_names and integ_outcome_vector_names consistent? make class TypeContext? related to lin
-        self.stan_model_context.target_integ_outcome_vector_names = target_simulated_vector_names
+        # TODO @Dashadower how to make target_sim_vector_names and integ_outcome_vector_names consistent? make class TypeContext? related to lin
+        self.stan_model_context.target_integ_outcome_vector_names = target_sim_vector_names
         self.stan_model_context.obs_integ_outcome_vector_names = [f'{name}_obs' for name in self.stan_model_context.target_integ_outcome_vector_names]
         self.driving_vector_names = driving_vector_names
         self.model_name = model_name
 
-    #TODO @dashadower for external refernce of target_simulated_vector_names
-    # is it better to use model.target_simulated_vector_names or define function for consistency?
+    #TODO @dashadower for external refernce of target_sim_vector_names
+    # is it better to use model.target_sim_vector_names or define function for consistency?
     def get_latent_vector_names(self):
         return [f'{target}' for target in self.stan_model_context.target_integ_outcome_vector_names]
 
@@ -211,9 +212,9 @@ def get_latent_obs_vector_names(self):
         return self.get_latent_vector_names() + self.get_obs_vector_names()
 
 
-    def update_setting(self, est_param_names: list, target_simulated_vector_names: list, driving_vector_names: list, model_name: str):
+    def update_setting(self, est_param_names: list, target_sim_vector_names: list, driving_vector_names: list, model_name: str):
         self.est_param_names = est_param_names
-        self.stan_model_context.target_integ_outcome_vector_names = target_simulated_vector_names
+        self.stan_model_context.target_integ_outcome_vector_names = target_sim_vector_names
         self.driving_vector_names = driving_vector_names
         self.model_name = model_name
         # if self.initial_time in self.integration_times:
@@ -386,7 +387,7 @@ def stanify_data2draws(self):
             f.write(StanModelBuilder(self.precision_context, self.stan_model_context).build_block(self.hier_est_param_names))
             f.write("\n")
 
-            f.write(Data2DrawsStanGQBuilder(self.precision_context,self.stan_model_context, self.vensim_model_context,).build_block())
+            f.write(Data2DrawsStanGQBuilder(self.precision_context,self.stan_model_context, self.vensim_model_context,).build_block(self.hier_est_param_names))
 
         stan_model = cmdstanpy.CmdStanModel(stan_file=stan_data2draws_path, cpp_options={'STAN_THREADS':'true'})
         return stan_model