Addes rastrigin_2d function

app.py

@@ -35,7 +35,7 @@ with st.container():
 # available_functions = ["testfunc_2d", "camel3_2d", "camel16_2d", "treccani_2d", "goldstein_2d", "branin_2d", 
 #                        "rosenbrock_2d", "ackley_2d"]
 
-available_functions = ["testfunc_2d", "camel3_2d", "camel16_2d", "treccani_2d", "goldstein_2d", "branin_2d", 
+available_functions = ["testfunc_2d", "rastrigin_2d", "camel3_2d", "camel16_2d", "treccani_2d", "goldstein_2d", "branin_2d", 
                        "rosenbrock_2d", "ackley_2d", "rosenbrock_10d", "griewank_10d", "rastrigin_10d"]
 
 st.sidebar.header('User Inputs')

src/functions.py

@@ -104,7 +104,15 @@ def rastrigin_10d(X):
     except:
         X = np.array([X])
     return (np.sum(X**2 - 10 * np.cos(2 * np.pi * X) + 10, axis=1)).reshape(-1, 1)
-    
+
+# 2d Rastrigin
+def rastrigin_2d(X):
+    try:
+        X.shape[1]
+    except:
+        X = np.array([X])
+    return (np.sum(X**2 - 10 * np.cos(2 * np.pi * X) + 10, axis=1)).reshape(-1, 1)
+
 # ------- END FUNCTIONS -------
 
 if __name__ == "__main__":
@@ -141,6 +149,9 @@ if __name__ == "__main__":
     elif f == rastrigin_10d:
         x1_min = x2_min = -5.12
         x1_max = x2_max = 5.12
+    elif f == rastrigin_2d:
+        x1_min = x2_min = -5.12
+        x1_max = x2_max = 5.12
     
 
     N = 100