def inv_map(self, sol: Solution) -> Solution:
if self.ind_state_idx is None:
self.ind_state_idx = np.shape(sol.y)[1] - 1
sol.t = sol.y[:, self.ind_state_idx]
sol.y = np.delete(sol.y, self.ind_state_idx, axis=1)
sol.lam_t = sol.lam[:, self.ind_state_idx]
sol.lam = np.delete(sol.lam, self.ind_state_idx, axis=1)
return sol
def map(self, sol: Solution) -> Solution:
if self.ind_state_idx is None:
self.ind_state_idx = np.shape(sol.y)[1]
if len(sol.lam_t) == 0:
sol.lam_t = np.zeros_like(sol.t)
sol.y = np.insert(sol.y, self.ind_state_idx, sol.t, axis=1)
sol.lam = np.insert(sol.lam, self.ind_state_idx, sol.lam_t, axis=1)
return sol
