Repository: pybamm-team/PyBaMM (python)

Issue:
[Bug]: IDAKLU SOLVER output_variables long simulation and state issues
### PyBaMM Version

24.11.2

### Python Version

3.12.4

### Describe the bug

IDAKLU SOLVER output_variables long simulation and state issues:

Hi! I have been running ageing simulations with CASADI for a while and wanted to use IDAKLU for CC ageing or calendar ageing as it is faster. With the new option for output_variables, the memory allocation is much smaller (x10), good for parallel or threaded computing. 

Now, two issues arise. Importantly, none of these errors occur with Casadi or IDAKLU (when output_variables are not specified)


**So basically using last_state in looped simulations and using .cycles to retrieve cycle specif data does not work too well with IDAKLU when output_variables is active :)**


**1)** When using cycles and selecting output_variables, one cannot use the previous way of working:
 for example like this: simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries
but instead must go with a work around like this: 
    **capacity_full_checkup = simulation_IDA3_checkup.solution.get_data_dict()['Throughput capacity [A.h]']
    cycles = simulation_IDA3_checkup.solution.get_data_dict()['Cycle']
    capacity_discharge_section= [d for d, c in zip(capacity_full_checkup, cycles) if c == 2]**   
I find this work around inconvenient and would prefer that what works with CASADI and IDAKLU (without output_Variables selected) works also with output_variables selected

**2)** when using cycling and for example using last states, I get an error with starting_solution
example:

    simulation_IDA3_ageing_cylic.solve(starting_solution=last_state)    
    last_state = simulation_IDA3_ageing_cylic.solution.last_state
    simulation_IDA3_checkup.solve(starting_solution=last_state)

this will give me an error if simulation_IDA3_ageing_cylic = pybamm.Simulation(... , solver=solver_IDAKLU2 , ....)
but it will not give me an error if  simulation_IDA3_ageing_cylic = pybamm.Simulation(... , solver=solver_IDAKLU1 , ....)
solver_IDAKLU1 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None)
solver_IDAKLU2 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None, output_variables=['Time [s]', 'Current [A]', "Voltage [V]" ], options=None)

for this last example I have a work around which is to use **starting_solution = simulation_IDA3_ageing_calendar.solution**, so using the whole solution instead of the last_state (which is I think more efficient and has a smaller memory footprint)

for 2) this means passing on a simulation object which, iteration by iteration, will grow in size. Using save_cycles in the simulation options only leads to more errors which I will not address here. all in all, I can 

Ideally, I would like to be able to activate output_variable in the IDAKLU solver, while retaining the capability of calling and using:

- starting_solution=last_state (in .solve step)
- simulation.solution.cycles[-1]['XXX'].entries (after .solve )
- save_cycles = XXX (in simulation step)

Am I getting this wrong? I.e., recongizning that it might not work just as if I swapped CASADI, are there other workarounds possibile? Let me know :) Thanks for the help 

Alessio Lodge, TNO - Imperial

### Steps to Reproduce

# Code is above. Use any experiment and run with this:


model = pybamm.lithium_ion.DFN()
solver_CasadiSolverSafe = pybamm.CasadiSolver(dt_max = 100,mode = "safe", rtol=1e-7 ,atol=1e-7)
solver_IDAKLU1 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None)
solver_IDAKLU2 = pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None, output_variables=['Time [s]', 'Current [A]', "Voltage [V]" ], options=None)
solver_IDAKLU3 =  pybamm.IDAKLUSolver(rtol=0.00001, atol=1e-07, root_method='casadi', root_tol=1e-07, extrap_tol=None, output_variables=['Time [s]', 'Current [A]', "Voltage [V]", 'Throughput capacity [A.h]', 'Discharge capacity [A.h]'], options=  {"dt_max": 1000.0})

# FOR CASADI OR IDAKLU WITHOUT OUTPUT VARS
simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries


# FOR  IDAKLU WITH OUTPUT VARS **1)**
    capacity_full_checkup = simulation_IDA3_checkup.solution.get_data_dict()['Throughput capacity [A.h]']
    cycles = simulation_IDA3_checkup.solution.get_data_dict()['Cycle']
    capacity_discharge_section= [d for d, c in zip(capacity_full_checkup, cycles) if c == 2]   
#I find this work around inconvenient and would prefer that what works with CASADI and IDAKLU (without output_Variables selected) works also with #output_variables selected


# FOR  IDAKLU WITH OUTPUT VARS **2)**
starting_solution = simulation_IDA3_ageing_calendar.solution #


### Relevant log output

```shell
for 1) (when I make the error happen, without work around)

  File "c:\Users\XXX\Desktop\Journal pape\Code\Python_ageing_code_working_in_IDAKLU.py", line 81, in get_EOL_cap
    print(simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries - min(simulation_IDA3_checkup.solution.cycles[-1]['Throughput capacity [A.h]'].entries)) #DOESNT WORK WITH IDAKLU, WHY???                                                                                                                                                                             WITH IDAKLU, WHY???
          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\lodgeaa\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 710, in __getitem__
    self.update(key)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 590, in update
    self._update_variable(variable)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 608, in _update_variable
    raise KeyError(
KeyError: "Cannot process variable 'Throughput capacity [A.h]' as it was not part of the solve. Please re-run the solve with `output_variables` set to include this variable."


for 2) (when I make the error happen, without work around)

Traceback (most recent call last):
  File "c:\Users\XXX\Desktop\Journal pape\Code\python_ageing_code_working_in_IDAKLU_with_output_variables", line 173, in <module>
    simulation_IDA3_checkup.solve(starting_solution=last_state)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\simulation.py", line 864, in solve
    self._solution = self._solution + cycle_solution
                     ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\solution.py", line 994, in __add__
    v: self._variables[v]._update(other._variables[v], new_sol)
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\processed_variable_computed.py", line 475, in _update
    new_var = self.__class__(bv, bvc, bvd, new_sol)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\processed_variable_computed.py", line 88, in __init__
    self.initialise_0D()
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybamm\solvers\processed_variable_computed.py", line 202, in initialise_0D
    self._xr_data_array = xr.DataArray(entries, coords=[("t", self.t_pts)])
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\xarray\core\dataarray.py", line 456, in __init__
    coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\xarray\core\dataarray.py", line 195, in _infer_coords_and_dims
    _check_coords_dims(shape, new_coords, dims_tuple)
  File "C:\Users\XXX\AppData\Local\Programs\Python\Python312\Lib\site-packages\xarray\core\dataarray.py", line 129, in _check_coords_dims
    raise ValueError(
ValueError: conflicting sizes for dimension 't': length 143 on the data but length 144 on coordinate 't'
```

Pull request description:
Bugfix: IDAKLU output variables with Experiments
# Description

Fixes misc bugs relating to the use of the `output_variables` option with the IDAKLU solver and Experiments.

Partly addresses #4743 

## Type of change

- [x] Bug fix (non-breaking change which fixes an issue)

# Key checklist:

- [ ] No style issues: `$ pre-commit run` (or `$ nox -s pre-commit`) (see [CONTRIBUTING.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CONTRIBUTING.md#installing-and-using-pre-commit) for how to set this up to run automatically when committing locally, in just two lines of code)
- [ ] All tests pass: `$ python -m pytest` (or `$ nox -s tests`)
- [ ] The documentation builds: `$ python -m pytest --doctest-plus src` (or `$ nox -s doctests`)

You can run integration tests, unit tests, and doctests together at once, using `$ nox -s quick`.

## Further checks:

- [ ] Code is commented, particularly in hard-to-understand areas
- [ ] Tests added that prove fix is effective or that feature works


Patch:
diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py
index 5c446ce87..4f17c60d9 100644
--- a/src/pybamm/solvers/solution.py
+++ b/src/pybamm/solvers/solution.py
@@ -1130,6 +1130,10 @@ def make_cycle_solution(
         sum_sols.all_yps,
         sum_sols.variables_returned,
     )
+
+    if sum_sols.variables_returned:
+        cycle_solution._variables = sum_sols._variables
+
     cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi
     cycle_solution._sub_solutions = sum_sols.sub_solutions
 


Test patch:
diff --git a/tests/integration/test_solvers/test_idaklu.py b/tests/integration/test_solvers/test_idaklu.py
index 12c785143..be92acd9d 100644
--- a/tests/integration/test_solvers/test_idaklu.py
+++ b/tests/integration/test_solvers/test_idaklu.py
@@ -133,23 +133,23 @@ class TestIDAKLUSolver:
         t_combined = np.concatenate((sol3.t, t_interp_dense))
         t_combined = np.unique(t_combined)
         t_combined.sort()
-        np.testing.assert_array_almost_equal(sol3.t, t_combined)
+        np.testing.assert_allclose(sol3.t, t_combined)
 
         # 4. sparse t_eval + sparse t_interp
         sol4 = solver.solve(
             model_disc, t_eval_sparse, t_interp=t_interp_sparse, inputs=inputs
         )
-        np.testing.assert_array_almost_equal(sol4.t, np.array([t0, tf]))
+        np.testing.assert_allclose(sol4.t, np.array([t0, tf]))
 
         sols = [sol1, sol2, sol3, sol4]
         for sol in sols:
             # test that y[0] = to true solution
             true_solution = a_value * sol.t
-            np.testing.assert_array_almost_equal(sol.y[0], true_solution)
+            np.testing.assert_allclose(sol.y[0], true_solution)
 
             # test that y[1:3] = to true solution
             true_solution = b_value * sol.t
-            np.testing.assert_array_almost_equal(sol.y[1:3], true_solution)
+            np.testing.assert_allclose(sol.y[1:3], true_solution)
 
     def test_with_experiments(self):
         summary_vars = []
@@ -195,10 +195,16 @@ class TestIDAKLUSolver:
         np.testing.assert_array_equal(
             sols[0]["Pressure [Pa]"].data, sols[1]["Pressure [Pa]"].data
         )
-        np.testing.assert_array_almost_equal(
+        np.testing.assert_allclose(
             sols[0]["Voltage [V]"].data, sols[1]["Voltage [V]"].data
         )
 
         # check summary variables are the same if output variables are specified
         for var in model.summary_variables:
             assert summary_vars[0][var] == summary_vars[1][var]
+
+        # check variables are accessible for each cycle
+        np.testing.assert_allclose(
+            sols[0].cycles[-1]["Current [A]"].data,
+            sols[1].cycles[-1]["Current [A]"].data,
+        )