openscvx 0.1.2__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

openscvx/propagation.py

@@ -2,6 +2,7 @@ import numpy as np
 
 from openscvx.config import Config
 from openscvx.integrators import solve_ivp_diffrax_prop
+from openscvx.backend.parameter import Parameter
 
 
 def prop_aug_dy(
@@ -15,7 +16,29 @@ def prop_aug_dy(
     state_dot: callable,
     dis_type: str,
     N: int,
+    *params
 ) -> np.ndarray:
+    """Compute the augmented dynamics for propagation.
+    
+    This function computes the time-scaled dynamics for propagating the system state,
+    taking into account the discretization type (ZOH or FOH) and time dilation.
+    
+    Args:
+        tau (float): Current normalized time in [0,1].
+        x (np.ndarray): Current state vector.
+        u_current (np.ndarray): Control input at current node.
+        u_next (np.ndarray): Control input at next node.
+        tau_init (float): Initial normalized time.
+        node (int): Current node index.
+        idx_s (int): Index of time dilation variable in control vector.
+        state_dot (callable): Function computing state derivatives.
+        dis_type (str): Discretization type ("ZOH" or "FOH").
+        N (int): Number of nodes in trajectory.
+        *params: Additional parameters passed to state_dot.
+        
+    Returns:
+        np.ndarray: Time-scaled state derivatives.
+    """
     x = x[None, :]
 
     if dis_type == "ZOH":
@@ -24,37 +47,69 @@ def prop_aug_dy(
         beta = (tau - tau_init) * N
     u = u_current + beta * (u_next - u_current)
 
-    return u[:, idx_s] * state_dot(x, u[:, :-1], node).squeeze()
-
-
-def get_propagation_solver(state_dot, params):
-    def propagation_solver(V0, tau_grid, u_cur, u_next, tau_init, node, idx_s):
+    return u[:, idx_s] * state_dot(x, u[:, :-1], node, *params).squeeze()
+
+def get_propagation_solver(state_dot, settings, param_map):
+    """Create a propagation solver function.
+    
+    This function creates a solver that propagates the system state using the
+    specified dynamics and settings.
+    
+    Args:
+        state_dot (callable): Function computing state derivatives.
+        settings: Configuration settings for propagation.
+        param_map (dict): Mapping of parameter names to values.
+        
+    Returns:
+        callable: A function that solves the propagation problem.
+    """
+    def propagation_solver(V0, tau_grid, u_cur, u_next, tau_init, node, idx_s, save_time, mask, *params):
+        param_map_update = dict(zip(param_map.keys(), params))
         return solve_ivp_diffrax_prop(
             f=prop_aug_dy,
-            tau_final=tau_grid[1],
-            y_0=V0,
+            tau_final=tau_grid[1],  # scalar
+            y_0=V0,                 # shape (n_states,)
             args=(
-                u_cur,
-                u_next,
-                tau_init,
-                node,
-                idx_s,
-                state_dot,
-                params.dis.dis_type,
-                params.scp.n,
+                u_cur,             # shape (1, n_controls)
+                u_next,            # shape (1, n_controls)
+                tau_init,          # shape (1, 1)
+                node,              # shape (1, 1)
+                idx_s,             # int
+                state_dot,         # function or array
+                settings.dis.dis_type,
+                settings.scp.n,
+                *param_map_update.items(),
+                # additional named parameters as **kwargs
             ),
-            tau_0=tau_grid[0],
+            tau_0=tau_grid[0],      # scalar
+            save_time=save_time,    # shape (MAX_TAU_LEN,)
+            mask=mask              # shape (MAX_TAU_LEN,), dtype=bool
         )
 
     return propagation_solver
 
 
-def s_to_t(u, params: Config):
-    t = [0]
+
+
+def s_to_t(x, u, params: Config):
+    """Convert normalized time s to real time t.
+    
+    This function converts the normalized time variable s to real time t
+    based on the discretization type and time dilation factors.
+    
+    Args:
+        x: State trajectory.
+        u: Control trajectory.
+        params (Config): Configuration settings.
+        
+    Returns:
+        list: List of real time points.
+    """
+    t = [x.guess[:,params.sim.idx_t][0]]
     tau = np.linspace(0, 1, params.scp.n)
     for k in range(1, params.scp.n):
-        s_kp = u[k - 1, -1]
-        s_k = u[k, -1]
+        s_kp = u.guess[k - 1, -1]
+        s_k = u.guess[k, -1]
         if params.dis.dis_type == "ZOH":
             t.append(t[k - 1] + (tau[k] - tau[k - 1]) * (s_kp))
         else:
@@ -62,17 +117,30 @@ def s_to_t(u, params: Config):
     return t
 
 
-def t_to_tau(u, t, u_nodal, t_nodal, params: Config):
-    u_lam = lambda new_t: np.array(
-        [np.interp(new_t, t_nodal, u[:, i]) for i in range(u.shape[1])]
-    ).T
+def t_to_tau(u, t, t_nodal, params: Config):
+    """Convert real time t to normalized time tau.
+    
+    This function converts real time t to normalized time tau and interpolates
+    the control inputs accordingly.
+    
+    Args:
+        u: Control trajectory.
+        t (np.ndarray): Real time points.
+        t_nodal (np.ndarray): Nodal time points.
+        params (Config): Configuration settings.
+        
+    Returns:
+        tuple: (tau, u_interp) where tau is normalized time and u_interp is interpolated controls.
+    """
+    u_guess = u.guess
+    u_lam = lambda new_t: np.array([np.interp(new_t, t_nodal, u_guess[:,i]) for i in range(u_guess.shape[1])]).T
     u = np.array([u_lam(t_i) for t_i in t])
 
     tau = np.zeros(len(t))
     tau_nodal = np.linspace(0, 1, params.scp.n)
     for k in range(1, len(t)):
         k_nodal = np.where(t_nodal < t[k])[0][-1]
-        s_kp = u_nodal[k_nodal, -1]
+        s_kp = u_guess[k_nodal, -1]
         tp = t_nodal[k_nodal]
         tau_p = tau_nodal[k_nodal]
 
@@ -84,52 +152,84 @@ def t_to_tau(u, t, u_nodal, t_nodal, params: Config):
     return tau, u
 
 
-def simulate_nonlinear_time(x_0, u, tau_vals, t, params, propagation_solver):
-    states = np.empty(
-        (x_0.shape[0], 0)
-    )  # Initialize states as a 2D array with shape (n, 0)
-
-    tau = np.linspace(0, 1, params.scp.n)
-
-    u_lam = lambda new_t: np.array(
-        [np.interp(new_t, t, u[:, i]) for i in range(u.shape[1])]
-    ).T
-
-    # Bin the tau_vals into with respect to the uniform tau grid, tau
+def simulate_nonlinear_time(params, x, u, tau_vals, t, settings, propagation_solver):
+    """Simulate the nonlinear system dynamics over time.
+    
+    This function simulates the system dynamics using the optimal control sequence
+    and returns the resulting state trajectory.
+    
+    Args:
+        params: System parameters.
+        x: State trajectory.
+        u: Control trajectory.
+        tau_vals (np.ndarray): Normalized time points for simulation.
+        t (np.ndarray): Real time points.
+        settings: Configuration settings.
+        propagation_solver (callable): Function for propagating the system state.
+        
+    Returns:
+        np.ndarray: Simulated state trajectory.
+    """
+    x_0 = settings.sim.x_prop.initial
+
+    n_segments = settings.scp.n - 1
+    n_states = x_0.shape[0]
+    n_tau = len(tau_vals)
+
+    params = params.items()
+    param_values = tuple([param.value for _, param in params])
+    
+    states = np.empty((n_states, n_tau))
+    tau = np.linspace(0, 1, settings.scp.n)
+
+    # Precompute control interpolation
+    u_interp = np.stack([
+        np.interp(t, t, u.guess[:, i]) for i in range(u.guess.shape[1])
+    ], axis=-1)
+
+    # Bin tau_vals into segments of tau
     tau_inds = np.digitize(tau_vals, tau) - 1
-    # Force the last indice to be in the same bin as the previous ones
-    tau_inds = np.where(tau_inds == params.scp.n - 1, params.scp.n - 2, tau_inds)
+    tau_inds = np.where(tau_inds == settings.scp.n - 1, settings.scp.n - 2, tau_inds)
 
     prev_count = 0
+    out_idx = 0
 
-    for k in range(params.scp.n - 1):
-        controls_current = np.squeeze(u_lam(t[k]))[None, :]
-        controls_next = np.squeeze(u_lam(t[k + 1]))[None, :]
+    for k in range(n_segments):
+        controls_current = u_interp[k][None, :]
+        controls_next = u_interp[k + 1][None, :]
 
-        # Create a mask
+        # Mask for tau_vals in current segment
         mask = (tau_inds >= k) & (tau_inds < k + 1)
-
         count = np.sum(mask)
 
-        # Use count to grab the first count number of elements
-        tau_cur = tau_vals[prev_count : prev_count + count]
+        tau_cur = tau_vals[prev_count:prev_count + count]
+        tau_cur = np.concatenate([tau_cur, np.array([tau[k + 1]])])  # Always include final point
+        count += 1
+
+        # Pad to fixed length
+        pad_len = settings.prp.max_tau_len - count
+        tau_cur_padded = np.pad(tau_cur, (0, pad_len), constant_values=tau[k + 1])
+        mask_padded = np.concatenate([np.ones(count), np.zeros(pad_len)]).astype(bool)
 
-        sol = propagation_solver(
+        # Call the solver with padded tau_cur and mask
+        sol = propagation_solver.call(
             x_0,
             (tau[k], tau[k + 1]),
             controls_current,
             controls_next,
             np.array([[tau[k]]]),
             np.array([[k]]),
-            params.sim.idx_s.stop,
+            settings.sim.idx_s.stop,
+            tau_cur_padded,
+            mask_padded,
+            *param_values
         )
 
-        x = sol.ys
-        for tau_i in tau_cur:
-            new_state = sol.evaluate(tau_i).reshape(-1, 1)  # Ensure new_state is 2D
-            states = np.concatenate([states, new_state], axis=1)
+        # Only store the valid portion (excluding the final point which becomes next x_0)
+        states[:, out_idx:out_idx + count - 1] = sol[:count - 1].T
+        out_idx += count - 1
+        x_0 = sol[count - 1]  # Last value used as next x_0
 
-        x_0 = x[-1]
-        prev_count += count
+        prev_count += (count - 1)
 
-    return states.T
+    return states.T

openscvx/ptr.py

@@ -4,55 +4,91 @@ import cvxpy as cp
 import pickle
 import time
 
+from openscvx.backend.parameter import Parameter
 from openscvx.config import Config
+from openscvx.results import OptimizationResults
 
 import warnings
 warnings.filterwarnings("ignore")
 
-def PTR_init(ocp: cp.Problem, discretization_solver: callable, params: Config):
-    if params.cvx.cvxpygen:
-        from solver.cpg_solver import cpg_solve
-        with open('solver/problem.pickle', 'rb') as f:
-            prob = pickle.load(f)
+def PTR_init(params, ocp: cp.Problem, discretization_solver: callable, settings: Config):
+    if settings.cvx.cvxpygen:
+        try:
+            from solver.cpg_solver import cpg_solve
+            with open('solver/problem.pickle', 'rb') as f:
+                prob = pickle.load(f)
+        except ImportError:
+            raise ImportError(
+                "cvxpygen solver not found. Make sure cvxpygen is installed and code generation has been run. "
+                "Install with: pip install openscvx[cvxpygen]"
+            )
     else:
         cpg_solve = None
 
+    if 'x_init' in ocp.param_dict:
+        ocp.param_dict['x_init'].value = settings.sim.x.initial
+    
+    if 'x_term' in ocp.param_dict:
+        ocp.param_dict['x_term'].value = settings.sim.x.final
+
     # Solve a dumb problem to intilize DPP and JAX jacobians
-    _ = PTR_subproblem(cpg_solve, params.sim.x_bar, params.sim.u_bar, discretization_solver, ocp, params)
+    _ = PTR_subproblem(params.items(), cpg_solve, settings.sim.x, settings.sim.u, discretization_solver, ocp, settings)
 
     return cpg_solve
 
-def PTR_main(params: Config, prob: cp.Problem, aug_dy: callable, cpg_solve, emitter_function) -> dict:
+def format_result(problem, converged: bool) -> OptimizationResults:
+    """Formats the final result as an OptimizationResults object from the problem's state."""
+    return OptimizationResults(
+        converged=converged,
+        t_final=problem.settings.sim.x.guess[:, problem.settings.sim.idx_t][-1],
+        u=problem.settings.sim.u,
+        x=problem.settings.sim.x,
+        x_history=problem.scp_trajs,
+        u_history=problem.scp_controls,
+        discretization_history=problem.scp_V_multi_shoot_traj,
+        J_tr_history=problem.scp_J_tr,
+        J_vb_history=problem.scp_J_vb,
+        J_vc_history=problem.scp_J_vc,
+    )
+
+def PTR_main(params, settings: Config, prob: cp.Problem, aug_dy: callable, cpg_solve, emitter_function) -> OptimizationResults:
     J_vb = 1E2
     J_vc = 1E2
     J_tr = 1E2
 
-    x_bar = params.sim.x_bar
-    u_bar = params.sim.u_bar
+    x = settings.sim.x
+    u = settings.sim.u
+
+    if 'x_init' in prob.param_dict:
+        prob.param_dict['x_init'].value = settings.sim.x.initial
+    
+    if 'x_term' in prob.param_dict:
+        prob.param_dict['x_term'].value = settings.sim.x.final
+
 
-    scp_trajs = [x_bar]
-    scp_controls = [u_bar]
+    scp_trajs = [x.guess]
+    scp_controls = [u.guess]
     V_multi_shoot_traj = []
 
     k = 1
 
-    while k <= params.scp.k_max and ((J_tr >= params.scp.ep_tr) or (J_vb >= params.scp.ep_vb) or (J_vc >= params.scp.ep_vc)):
-        x, u, t, J_total, J_vb_vec, J_vc_vec, J_tr_vec, prob_stat, V_multi_shoot, subprop_time, dis_time = PTR_subproblem(cpg_solve, x_bar, u_bar, aug_dy, prob, params)
+    while k <= settings.scp.k_max and ((J_tr >= settings.scp.ep_tr) or (J_vb >= settings.scp.ep_vb) or (J_vc >= settings.scp.ep_vc)):
+        x_sol, u_sol, cost, J_total, J_vb_vec, J_vc_vec, J_tr_vec, prob_stat, V_multi_shoot, subprop_time, dis_time = PTR_subproblem(params.items(), cpg_solve, x, u, aug_dy, prob, settings)
 
         V_multi_shoot_traj.append(V_multi_shoot)
 
-        x_bar = x
-        u_bar = u
+        x.guess = x_sol
+        u.guess = u_sol
 
         J_tr = np.sum(np.array(J_tr_vec))
         J_vb = np.sum(np.array(J_vb_vec))
         J_vc = np.sum(np.array(J_vc_vec))
-        scp_trajs.append(x)
-        scp_controls.append(u)
+        scp_trajs.append(x.guess)
+        scp_controls.append(u.guess)
 
-        params.scp.w_tr = min(params.scp.w_tr * params.scp.w_tr_adapt, params.scp.w_tr_max)
-        if k > params.scp.cost_drop:
-            params.scp.lam_cost = params.scp.lam_cost * params.scp.cost_relax
+        settings.scp.w_tr = min(settings.scp.w_tr * settings.scp.w_tr_adapt, settings.scp.w_tr_max)
+        if k > settings.scp.cost_drop:
+            settings.scp.lam_cost = settings.scp.lam_cost * settings.scp.cost_relax
 
         emitter_function(
             {
@@ -63,34 +99,37 @@ def PTR_main(params: Config, prob: cp.Problem, aug_dy: callable, cpg_solve, emit
                 "J_tr": J_tr,
                 "J_vb": J_vb,
                 "J_vc": J_vc,
-                "cost": t[-1],
+                "cost": cost[-1],
                 "prob_stat": prob_stat,
             }
         )
 
         k += 1
 
-    result = dict(
-        converged = k <= params.scp.k_max,
-        t_final = x[:,params.sim.idx_t][-1],
-        u = u,
-        x = x,
-        x_history = scp_trajs,
-        u_history = scp_controls,
-        discretization_history = V_multi_shoot_traj,
-        J_tr_history = J_tr_vec,
-        J_vb_history = J_vb_vec,
-        J_vc_history = J_vc_vec,
+    result = OptimizationResults(
+        converged=k <= settings.scp.k_max,
+        t_final=x.guess[:,settings.sim.idx_t][-1],
+        u=u,
+        x=x,
+        x_history=scp_trajs,
+        u_history=scp_controls,
+        discretization_history=V_multi_shoot_traj,
+        J_tr_history=J_tr_vec,
+        J_vb_history=J_vb_vec,
+        J_vc_history=J_vc_vec,
     )
+    
     return result
 
-
-def PTR_subproblem(cpg_solve, x_bar, u_bar, aug_dy, prob, params: Config):
-    prob.param_dict['x_bar'].value = x_bar
-    prob.param_dict['u_bar'].value = u_bar
+def PTR_subproblem(params, cpg_solve, x, u, aug_dy, prob, settings: Config):
+    prob.param_dict['x_bar'].value = x.guess
+    prob.param_dict['u_bar'].value = u.guess
     
+    # Make a tuple from list of parameter values
+    param_values = tuple([param.value for _, param in params])
+
     t0 = time.time()
-    A_bar, B_bar, C_bar, z_bar, V_multi_shoot = aug_dy(x_bar, u_bar.astype(float))
+    A_bar, B_bar, C_bar, z_bar, V_multi_shoot = aug_dy.call(x.guess, u.guess.astype(float), *param_values)
 
     prob.param_dict['A_d'].value = A_bar.__array__()
     prob.param_dict['B_d'].value = B_bar.__array__()
@@ -98,52 +137,52 @@ def PTR_subproblem(cpg_solve, x_bar, u_bar, aug_dy, prob, params: Config):
     prob.param_dict['z_d'].value = z_bar.__array__()
     dis_time = time.time() - t0
 
-    if params.sim.constraints_nodal:
-        for g_id, constraint in enumerate(params.sim.constraints_nodal):
+    if settings.sim.constraints_nodal:
+        for g_id, constraint in enumerate(settings.sim.constraints_nodal):
             if not constraint.convex:
-                prob.param_dict['g_' + str(g_id)].value = np.asarray(constraint.g(x_bar, u_bar))
-                prob.param_dict['grad_g_x_' + str(g_id)].value = np.asarray(constraint.grad_g_x(x_bar, u_bar))
-                prob.param_dict['grad_g_u_' + str(g_id)].value = np.asarray(constraint.grad_g_u(x_bar, u_bar))
+                prob.param_dict['g_' + str(g_id)].value = np.asarray(constraint.g(x.guess, u.guess))
+                prob.param_dict['grad_g_x_' + str(g_id)].value = np.asarray(constraint.grad_g_x(x.guess, u.guess))
+                prob.param_dict['grad_g_u_' + str(g_id)].value = np.asarray(constraint.grad_g_u(x.guess, u.guess))
     
-    prob.param_dict['w_tr'].value = params.scp.w_tr
-    prob.param_dict['lam_cost'].value = params.scp.lam_cost
+    prob.param_dict['w_tr'].value = settings.scp.w_tr
+    prob.param_dict['lam_cost'].value = settings.scp.lam_cost
 
-    if params.cvx.cvxpygen:
+    if settings.cvx.cvxpygen:
         t0 = time.time()
         prob.register_solve('CPG', cpg_solve)
-        prob.solve(method = 'CPG', **params.cvx.solver_args)
+        prob.solve(method = 'CPG', **settings.cvx.solver_args)
         subprop_time = time.time() - t0
     else:
         t0 = time.time()
-        prob.solve(solver = params.cvx.solver, enforce_dpp = True, **params.cvx.solver_args)
+        prob.solve(solver = settings.cvx.solver, **settings.cvx.solver_args)
         subprop_time = time.time() - t0
 
-    x = (params.sim.S_x @ prob.var_dict['x'].value.T + np.expand_dims(params.sim.c_x, axis = 1)).T
-    u = (params.sim.S_u @ prob.var_dict['u'].value.T + np.expand_dims(params.sim.c_u, axis = 1)).T
+    x_new_guess = (settings.sim.S_x @ prob.var_dict['x'].value.T + np.expand_dims(settings.sim.c_x, axis = 1)).T
+    u_new_guess = (settings.sim.S_u @ prob.var_dict['u'].value.T + np.expand_dims(settings.sim.c_u, axis = 1)).T
 
     i = 0
     costs = [0]
-    for type in params.sim.final_state.type:
+    for type in x.final_type:
         if type == 'Minimize':
-            costs += x[:,i]
+            costs += x_new_guess[:,i]
         if type == 'Maximize':
-            costs -= x[:,i]
+            costs -= x_new_guess[:,i]
         i += 1
 
     # Create the block diagonal matrix using jax.numpy.block
     inv_block_diag = np.block([
-        [params.sim.inv_S_x, np.zeros((params.sim.inv_S_x.shape[0], params.sim.inv_S_u.shape[1]))],
-        [np.zeros((params.sim.inv_S_u.shape[0], params.sim.inv_S_x.shape[1])), params.sim.inv_S_u]
+        [settings.sim.inv_S_x, np.zeros((settings.sim.inv_S_x.shape[0], settings.sim.inv_S_u.shape[1]))],
+        [np.zeros((settings.sim.inv_S_u.shape[0], settings.sim.inv_S_x.shape[1])), settings.sim.inv_S_u]
     ])
 
     # Calculate J_tr_vec using the JAX-compatible block diagonal matrix
-    J_tr_vec = la.norm(inv_block_diag @ np.hstack((x - x_bar, u - u_bar)).T, axis=0)**2
+    J_tr_vec = la.norm(inv_block_diag @ np.hstack((x_new_guess - x.guess, u_new_guess - u.guess)).T, axis=0)**2
     J_vc_vec = np.sum(np.abs(prob.var_dict['nu'].value), axis = 1)
     
     id_ncvx = 0
     J_vb_vec = 0
-    for constraint in params.sim.constraints_nodal:
+    for constraint in settings.sim.constraints_nodal:
         if constraint.convex == False:
             J_vb_vec += np.maximum(0, prob.var_dict['nu_vb_' + str(id_ncvx)].value)
             id_ncvx += 1
-    return x, u, costs, prob.value, J_vb_vec, J_vc_vec, J_tr_vec, prob.status, V_multi_shoot, subprop_time, dis_time
+    return x_new_guess, u_new_guess, costs, prob.value, J_vb_vec, J_vc_vec, J_tr_vec, prob.status, V_multi_shoot, subprop_time, dis_time

openscvx/results.py

@@ -0,0 +1,153 @@
+from dataclasses import dataclass, field
+from typing import List, Optional, Any, Dict
+import numpy as np
+
+from openscvx.backend.state import State
+from openscvx.backend.control import Control
+
+
+@dataclass
+class OptimizationResults:
+    """
+    Class to hold optimization results from the SCP solver.
+    
+    This class replaces the dictionary-based results structure with a more
+    structured and type-safe approach.
+    """
+    
+    # Core optimization results
+    converged: bool
+    t_final: float
+    u: Control
+    x: State
+    
+    # History of SCP iterations
+    x_history: List[np.ndarray] = field(default_factory=list)
+    u_history: List[np.ndarray] = field(default_factory=list)
+    discretization_history: List[np.ndarray] = field(default_factory=list)
+    J_tr_history: List[np.ndarray] = field(default_factory=list)
+    J_vb_history: List[np.ndarray] = field(default_factory=list)
+    J_vc_history: List[np.ndarray] = field(default_factory=list)
+    
+    # Post-processing results (added by propagate_trajectory_results)
+    t_full: Optional[np.ndarray] = None
+    x_full: Optional[np.ndarray] = None
+    u_full: Optional[np.ndarray] = None
+    cost: Optional[float] = None
+    ctcs_violation: Optional[np.ndarray] = None
+    
+    # Additional plotting/application data (added by user)
+    plotting_data: Dict[str, Any] = field(default_factory=dict)
+    
+    def __post_init__(self):
+        """Initialize the results object."""
+        pass
+    
+    def update_plotting_data(self, **kwargs):
+        """
+        Update the plotting data with additional information.
+        
+        Args:
+            **kwargs: Key-value pairs to add to plotting_data
+        """
+        self.plotting_data.update(kwargs)
+    
+    def get(self, key: str, default: Any = None) -> Any:
+        """
+        Get a value from the results, similar to dict.get().
+        
+        Args:
+            key: The key to look up
+            default: Default value if key is not found
+            
+        Returns:
+            The value associated with the key, or default if not found
+        """
+        # Check if it's a direct attribute
+        if hasattr(self, key):
+            return getattr(self, key)
+        
+        # Check if it's in plotting_data
+        if key in self.plotting_data:
+            return self.plotting_data[key]
+        
+        return default
+    
+    def __getitem__(self, key: str) -> Any:
+        """
+        Allow dictionary-style access to results.
+        
+        Args:
+            key: The key to look up
+            
+        Returns:
+            The value associated with the key
+            
+        Raises:
+            KeyError: If key is not found
+        """
+        # Check if it's a direct attribute
+        if hasattr(self, key):
+            return getattr(self, key)
+        
+        # Check if it's in plotting_data
+        if key in self.plotting_data:
+            return self.plotting_data[key]
+        
+        raise KeyError(f"Key '{key}' not found in results")
+    
+    def __setitem__(self, key: str, value: Any):
+        """
+        Allow dictionary-style assignment to results.
+        
+        Args:
+            key: The key to set
+            value: The value to assign
+        """
+        # Check if it's a direct attribute
+        if hasattr(self, key):
+            setattr(self, key, value)
+        else:
+            # Store in plotting_data
+            self.plotting_data[key] = value
+    
+    def __contains__(self, key: str) -> bool:
+        """
+        Check if a key exists in the results.
+        
+        Args:
+            key: The key to check
+            
+        Returns:
+            True if key exists, False otherwise
+        """
+        return hasattr(self, key) or key in self.plotting_data
+    
+    def update(self, other: Dict[str, Any]):
+        """
+        Update the results with additional data from a dictionary.
+        
+        Args:
+            other: Dictionary containing additional data
+        """
+        for key, value in other.items():
+            self[key] = value
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the results to a dictionary for backward compatibility.
+        
+        Returns:
+            Dictionary representation of the results
+        """
+        result_dict = {}
+        
+        # Add all direct attributes
+        for attr_name in self.__dataclass_fields__:
+            if attr_name != 'plotting_data':
+                result_dict[attr_name] = getattr(self, attr_name)
+        
+        # Add plotting data
+        result_dict.update(self.plotting_data)
+        
+        return result_dict