new models

ssms/basic_simulators/simulator.py

@@ -327,6 +327,8 @@ def check_if_z_gt_a(z, a):
             check_lba_drifts_sum(theta["v_RL"])
             check_lba_drifts_sum(theta["v_WM"])
             check_if_z_gt_a(theta["z"], theta["a"])
+        elif model in ["rlwm_lba_race_v2"]:
+            check_if_z_gt_a(theta["z"], theta["a"])
         elif model in ["lba_angle_3_v2", "rlwm_lba_pw_v1"]:
             check_if_z_gt_a(theta["z"], theta["a"])
 
@@ -541,7 +543,7 @@ def simulator(
         theta["z"] = np.expand_dims(theta["z"], axis=1)
         theta["theta"] = np.expand_dims(theta["theta"], axis=1)
 
-    if model == "rlwm_lba_race_v1":
+    if model in ["rlwm_lba_race_v1", "rlwm_lba_race_v2"]:
         sim_param_dict["sd"] = noise_dict["lba_based_models"]
         theta["v_RL"] = np.column_stack(
             [theta["v_RL_0"], theta["v_RL_1"], theta["v_RL_2"]]
@@ -551,7 +553,7 @@ def simulator(
         )
         theta["a"] = np.expand_dims(theta["a"], axis=1)
         theta["z"] = np.expand_dims(theta["z"], axis=1)
-    
+
     if model == "rlwm_lba_pw_v1":
         sim_param_dict["sd"] = noise_dict["lba_based_models"]
         theta["v_RL"] = np.column_stack(

ssms/config/config.py

@@ -310,7 +310,7 @@
         "simulator": cssm.lba_angle,
     },
     "rlwm_lba_pw_v1": {
-        "name": "rlwm_lba_pw_v1", 
+        "name": "rlwm_lba_pw_v1",
         "params": [
             "v_RL_0",
             "v_RL_1",
@@ -320,7 +320,7 @@
             "v_WM_2",
             "a",
             "z",
-            "t_WM"
+            "t_WM",
         ],
         "param_bounds": [
             [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.05, 0.0, 0.01],

ssms/dataset_generators/lan_mlp.py

@@ -259,7 +259,7 @@ def parameter_transform_for_data_gen(self, theta):
                 tmp = theta[3]
                 theta[3] = theta[4]
                 theta[4] = tmp
-        
+
         if self.model_config["name"] == "rlwm_lba_pw_v1":
             # ensure that a is always greater than z
             if theta[6] <= theta[7]: