Update app.py

app.py

@@ -124,6 +124,10 @@ if st.button("Set Up"):
         grad = sp.diff(expr_final, x)
         gradient_func = sp.lambdify(x, grad, "numpy")
         
+        # Store the function and gradient in session state
+        st.session_state.func = func
+        st.session_state.gradient_func = gradient_func
+        
         # Initialize session state for points and steps
         if 'points' not in st.session_state:
             st.session_state.points = [start_point]
@@ -136,12 +140,12 @@ if st.button("Set Up"):
         st.error(f"Error setting up function: {e}")
 
 # Gradient Descent Iteration button
-if 'points' in st.session_state and 'step' in st.session_state:
+if 'func' in st.session_state and 'gradient_func' in st.session_state:
     if st.button("Next Iteration"):
         try:
             # Get the current point and gradient value
             x_old = float(st.session_state.points[-1])
-            grad_val = gradient_func(x_old)
+            grad_val = st.session_state.gradient_func(x_old)
             x_new = x_old - learning_rate * grad_val
             
             # Append the new point to the list of points
@@ -154,15 +158,15 @@ if 'points' in st.session_state and 'step' in st.session_state:
             st.error(f"Error in iteration: {e}")
 
 # Creating the plot
-if 'points' in st.session_state and len(st.session_state.points) > 0:
+if 'func' in st.session_state and len(st.session_state.points) > 0:
     try:
         # Create x-values for plotting the function
         x_val = np.linspace(-6, 6, 500)
-        y_val = func(x_val)
+        y_val = st.session_state.func(x_val)
         
         # Get the points visited by gradient descent
         iter_points = np.array(st.session_state.points)
-        iter_y = func(iter_points)
+        iter_y = st.session_state.func(iter_points)
 
         # Plot the function and the gradient descent path
         trace1 = go.Scatter(x=x_val, y=y_val, mode="lines", name="Function", line=dict(color="blue"))
@@ -186,4 +190,3 @@ if 'points' in st.session_state and len(st.session_state.points) > 0:
 
     except Exception as e:
         st.error(f"Plot error: {e}")
-