Repository: pybamm-team/PyBaMM (python)

Issue:
[Bug]: Multi-particle model does not work with stoichiometry-dependent solid diffusivity 
### PyBaMM Version

23.9

### Python Version

3.10

### Describe the bug

I am getting an error when running a multi-particle DFN with a stoichiometry-dependent solid particle diffusivity:
`UnboundLocalError: local variable 'matrix' referenced before assignment` 

It seems that the `integration_dimension` is set to tertiary when there is stoichiometry-dependent solid particle diffusivity, but tertiary integration_dimension is [not yet implemented here](https://github.com/pybamm-team/PyBaMM/blob/6851496ff1a9cc3ec1573033cec28a2353857fbf/pybamm/spatial_methods/finite_volume.py#L216). 

### Steps to Reproduce

Create a MP-DFN model

```
options = {
        "particle size": "distribution"
}

model = pybamm.lithium_ion.DFN(options = options)

where the solid diffusivity functional form is sto dependent
for example:
def particle_diffusivity(sto, T):

    D_ref = 1 * 10 ** (-13)
    E_D_s = 18550
    arrhenius = np.exp(E_D_s / pybamm.constants.R * (1 / 298.15 - 1 / T))

    return D_ref * arrhenius * (sto + 0.2) * 10

params = pybamm.ParameterValues("Marquis2019")
params = pybamm.get_size_distribution_parameters(params)
params.update({
        "Positive electrode diffusivity [m2.s-1]": particle_diffusivity
})

```

### Relevant log output

```shell
`2023-11-21 11:58:27.134 - [INFO] discretisation.remove_independent_variables_from_rhs(1134): removing variable Discharge capacity [A.h] from rhs
2023-11-21 11:58:27.137 - [INFO] discretisation.remove_independent_variables_from_rhs(1134): removing variable Throughput capacity [A.h] from rhs
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:728, in Discretisation.process_symbol(self, symbol)
    727 try:
--> 728     return self._discretised_symbols[symbol]
    729 except KeyError:

KeyError: Addition(-0x20be135ef0fe3268, +, children=['2.0000000000000003e-13', 'x-average(1e-12 * (maximum(minimum(1.9524414269595433e-05 * Positive particle concentration distribution [mol.m-3], 0.9999999999), 1e-10)))'], domains={'primary': ['positive particle'], 'secondary': ['positive particle size'], 'tertiary': ['current collector']})

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:728, in Discretisation.process_symbol(self, symbol)
    727 try:
--> 728     return self._discretised_symbols[symbol]
    729 except KeyError:

KeyError: XAverage(-0x67124ee27e56970a, x-average, children=['1e-12 * (maximum(minimum(1.9524414269595433e-05 * Positive particle concentration distribution [mol.m-3], 0.9999999999), 1e-10))'], domains={'primary': ['positive particle'], 'secondary': ['positive particle size'], 'tertiary': ['current collector']})

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:728, in Discretisation.process_symbol(self, symbol)
    727 try:
--> 728     return self._discretised_symbols[symbol]
    729 except KeyError:

KeyError: Multiplication(-0x97cfd4073097872, *, children=["1.0 / integral dx_p ['positive particle'](broadcast(1.0))", "integral dx_p ['positive particle'](1e-12 * (maximum(minimum(1.9524414269595433e-05 * Positive particle concentration distribution [mol.m-3], 0.9999999999), 1e-10)))"], domains={'primary': ['positive particle'], 'secondary': ['positive particle size'], 'tertiary': ['current collector']})

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:728, in Discretisation.process_symbol(self, symbol)
    727 try:
--> 728     return self._discretised_symbols[symbol]
    729 except KeyError:

KeyError: Division(-0x3f2cb559d729ce0e, /, children=['1.0', "integral dx_p ['positive particle'](broadcast(1.0))"], domains={'primary': ['positive particle'], 'secondary': ['positive particle size'], 'tertiary': ['current collector']})

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:728, in Discretisation.process_symbol(self, symbol)
    727 try:
--> 728     return self._discretised_symbols[symbol]
    729 except KeyError:

KeyError: Integral(-0x3d067ca2936c7ceb, integral dx_p ['positive particle'], children=['broadcast(1.0)'], domains={'primary': ['positive particle'], 'secondary': ['positive particle size'], 'tertiary': ['current collector']})

During handling of the above exception, another exception occurred:

UnboundLocalError                         Traceback (most recent call last)
Cell In[86], line 20
     14 for model in models:
     15     sim = pybamm.Simulation(model, 
     16                             parameter_values=params,
     17 #                             experiment=experiment,
     18                            )
---> 20     sol = sim.solve(
     21         [0, 10],
     22 #         initial_soc="4 V"
     23     )
     24     sims.append(sol)
     25 sol.plot()

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/simulation.py:546, in Simulation.solve(self, t_eval, solver, check_model, save_at_cycles, calc_esoh, starting_solution, initial_soc, callbacks, showprogress, **kwargs)
    544 if self.operating_mode in ["without experiment", "drive cycle"]:
    545     print(build_inputs)
--> 546     self.build(check_model=check_model, initial_soc=initial_soc, inputs=build_inputs)
    547     if save_at_cycles is not None:
    548         raise ValueError(
    549             "'save_at_cycles' option can only be used if simulating an "
    550             "Experiment "
    551         )

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/simulation.py:427, in Simulation.build(self, check_model, initial_soc, inputs)
    425 self._mesh = pybamm.Mesh(self._geometry, self._submesh_types, self._var_pts)
    426 self._disc = pybamm.Discretisation(self._mesh, self._spatial_methods)
--> 427 self._built_model = self._disc.process_model(
    428     self._model_with_set_params, inplace=False, check_model=check_model
    429 )
    430 # rebuilt model so clear solver setup
    431 self._solver._model_set_up = {}

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:214, in Discretisation.process_model(self, model, inplace, check_model, remove_independent_variables_from_rhs)
    210 # Discretise variables (applying boundary conditions)
    211 # Note that we **do not** discretise the keys of model.rhs,
    212 # model.initial_conditions and model.boundary_conditions
    213 pybamm.logger.verbose("Discretise variables for {}".format(model.name))
--> 214 model_disc.variables = self.process_dict(model.variables)
    216 # Process parabolic and elliptic equations
    217 pybamm.logger.verbose("Discretise model equations for {}".format(model.name))

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:698, in Discretisation.process_dict(self, var_eqn_dict, ics)
    695         eqn = pybamm.FullBroadcast(eqn, broadcast_domains=eqn_key.domains)
    697 pybamm.logger.debug("Discretise {!r}".format(eqn_key))
--> 698 processed_eqn = self.process_symbol(eqn)
    699 # Calculate scale if the key has a scale
    700 scale = getattr(eqn_key, "scale", 1)

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:730, in Discretisation.process_symbol(self, symbol)
    728     return self._discretised_symbols[symbol]
    729 except KeyError:
--> 730     discretised_symbol = self._process_symbol(symbol)
    731     self._discretised_symbols[symbol] = discretised_symbol
    732     discretised_symbol.test_shape()

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:766, in Discretisation._process_symbol(self, symbol)
    764 left, right = symbol.children
    765 disc_left = self.process_symbol(left)
--> 766 disc_right = self.process_symbol(right)
    767 if symbol.domain == []:
    768     return pybamm.simplify_if_constant(
    769         symbol._binary_new_copy(disc_left, disc_right)
    770     )

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:730, in Discretisation.process_symbol(self, symbol)
    728     return self._discretised_symbols[symbol]
    729 except KeyError:
--> 730     discretised_symbol = self._process_symbol(symbol)
    731     self._discretised_symbols[symbol] = discretised_symbol
    732     discretised_symbol.test_shape()

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:789, in Discretisation._process_symbol(self, symbol)
    787         v = pybamm.ones_like(child)
    788         average = pybamm.Integral(child, x) / pybamm.Integral(v, x)
--> 789     return self.process_symbol(average)
    791 elif isinstance(symbol, pybamm.UnaryOperator):
    792     child = symbol.child

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:730, in Discretisation.process_symbol(self, symbol)
    728     return self._discretised_symbols[symbol]
    729 except KeyError:
--> 730     discretised_symbol = self._process_symbol(symbol)
    731     self._discretised_symbols[symbol] = discretised_symbol
    732     discretised_symbol.test_shape()

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:765, in Discretisation._process_symbol(self, symbol)
    762 if isinstance(symbol, pybamm.BinaryOperator):
    763     # Pre-process children
    764     left, right = symbol.children
--> 765     disc_left = self.process_symbol(left)
    766     disc_right = self.process_symbol(right)
    767     if symbol.domain == []:

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:730, in Discretisation.process_symbol(self, symbol)
    728     return self._discretised_symbols[symbol]
    729 except KeyError:
--> 730     discretised_symbol = self._process_symbol(symbol)
    731     self._discretised_symbols[symbol] = discretised_symbol
    732     discretised_symbol.test_shape()

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:766, in Discretisation._process_symbol(self, symbol)
    764 left, right = symbol.children
    765 disc_left = self.process_symbol(left)
--> 766 disc_right = self.process_symbol(right)
    767 if symbol.domain == []:
    768     return pybamm.simplify_if_constant(
    769         symbol._binary_new_copy(disc_left, disc_right)
    770     )

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:730, in Discretisation.process_symbol(self, symbol)
    728     return self._discretised_symbols[symbol]
    729 except KeyError:
--> 730     discretised_symbol = self._process_symbol(symbol)
    731     self._discretised_symbols[symbol] = discretised_symbol
    732     discretised_symbol.test_shape()

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/discretisations/discretisation.py:831, in Discretisation._process_symbol(self, symbol)
    827 elif isinstance(symbol, pybamm.Integral):
    828     integral_spatial_method = self.spatial_methods[
    829         symbol.integration_variable[0].domain[0]
    830     ]
--> 831     out = integral_spatial_method.integral(
    832         child, disc_child, symbol._integration_dimension
    833     )
    834     out.copy_domains(symbol)
    835     return out

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/spatial_methods/finite_volume.py:209, in FiniteVolume.integral(self, child, discretised_child, integration_dimension)
    207 def integral(self, child, discretised_child, integration_dimension):
    208     """Vector-vector dot product to implement the integral operator."""
--> 209     integration_vector = self.definite_integral_matrix(
    210         child, integration_dimension=integration_dimension
    211     )
    212     out = integration_vector @ discretised_child
    214     return out

File ~/PycharmProjects/PyBaMM-AB/PyBaMM/pybamm/spatial_methods/finite_volume.py:308, in FiniteVolume.definite_integral_matrix(self, child, vector_type, integration_dimension)
    301     matrix = kron(eye(third_dim_repeats), int_matrix)
    302 # print(domain, integration_dimension)
    303 # generate full matrix from the submatrix
    304 # Convert to csr_matrix so that we can take the index (row-slicing), which is
    305 # not supported by the default kron format
    306 # Note that this makes column-slicing inefficient, but this should not be an
    307 # issue
--> 308 return pybamm.Matrix(csr_matrix(matrix))

UnboundLocalError: local variable 'matrix' referenced before assignment`
```


Pull request description:
Issue 3549 psd diffusivity bug
# Description
Fixes a bug when using stoichiometry-dependent diffusivity with the DFN model with a particle size distribution

Fixes #3549 

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/develop/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

- [ ] New feature (non-breaking change which adds functionality)
- [ ] Optimization (back-end change that speeds up the code)
- [ ] 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/CHANGELOG.md b/CHANGELOG.md
index 388d66d89..8fcfb700e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,7 @@ package to install PyBaMM with only the required dependencies. ([conda-forge/pyb
 
 ## Bug fixes
 
+- Fixed bug when using stoichiometry-dependent diffusivity with the DFN model with a particle size distribution. ([#4726](https://github.com/pybamm-team/PyBaMM/pull/4726))
 - Remove internal use of deprecated `set_parameters` function in the `Simulation` class which caused warnings. ([#4638](https://github.com/pybamm-team/PyBaMM/pull/4638))
 - Provide default value for `Symbol.mesh` attribute to avoid errors when adding variables after discretisation. ([#4644](https://github.com/pybamm-team/PyBaMM/pull/4644))
 
diff --git a/src/pybamm/models/submodels/particle/base_particle.py b/src/pybamm/models/submodels/particle/base_particle.py
index e47039b97..80c09d7ac 100644
--- a/src/pybamm/models/submodels/particle/base_particle.py
+++ b/src/pybamm/models/submodels/particle/base_particle.py
@@ -166,7 +166,7 @@ class BaseParticle(pybamm.BaseSubModel):
 
         variables = {f"{Domain} {phase_name}particle flux [mol.m-2.s-1]": N_s}
 
-        if isinstance(N_s, pybamm.Broadcast):
+        if isinstance(N_s, pybamm.SecondaryBroadcast):
             N_s_xav = pybamm.x_average(N_s)
             variables.update(
                 {
@@ -403,29 +403,29 @@ class BaseParticle(pybamm.BaseSubModel):
 
         if [f"{domain} electrode"] in N_s.domains.values():
             # N_s depends on x
-
             N_s_distribution = N_s
-            # x-av the *tertiary* domain
-            # NOTE: not yet implemented. Fill with zeros instead
-            N_s_xav_distribution = pybamm.FullBroadcast(
-                0,
-                [f"{domain} {phase_name}particle"],
-                {
-                    "secondary": f"{domain} {phase_name}particle size",
-                    "tertiary": "current collector",
-                },
-            )
+            if isinstance(N_s, pybamm.TertiaryBroadcast):
+                N_s_xav_distribution = N_s.orphans[0]
+            else:
+                # can't average variables that evaluate on edges
+                N_s_xav_distribution = None
         else:
             N_s_xav_distribution = N_s
             N_s_distribution = pybamm.TertiaryBroadcast(N_s, [f"{domain} electrode"])
 
         variables = {
-            f"X-averaged {domain} {phase_name}particle flux "
-            "distribution [mol.m-2.s-1]": N_s_xav_distribution,
             f"{Domain} {phase_name}particle flux "
             "distribution [mol.m-2.s-1]": N_s_distribution,
         }
 
+        if N_s_xav_distribution is not None:
+            variables.update(
+                {
+                    f"X-averaged {domain} {phase_name}particle flux "
+                    "distribution [mol.m-2.s-1]": N_s_xav_distribution,
+                }
+            )
+
         return variables
 
     def _get_standard_diffusivity_variables(self, D_eff):
diff --git a/src/pybamm/models/submodels/particle/fickian_diffusion.py b/src/pybamm/models/submodels/particle/fickian_diffusion.py
index 5ae5d3a4d..cd67d6da9 100644
--- a/src/pybamm/models/submodels/particle/fickian_diffusion.py
+++ b/src/pybamm/models/submodels/particle/fickian_diffusion.py
@@ -130,7 +130,6 @@ class FickianDiffusion(BaseParticle):
     def get_coupled_variables(self, variables):
         domain, Domain = self.domain_Domain
         phase_name = self.phase_name
-        Phase_prefactor = self.phase_param.phase_prefactor
 
         if self.size_distribution is False:
             if self.x_average is False:
@@ -213,6 +212,14 @@ class FickianDiffusion(BaseParticle):
             }
         )
 
+        if self.x_average is True:
+            if self.size_distribution is True:
+                D_eff = pybamm.TertiaryBroadcast(D_eff, [f"{domain} electrode"])
+                N_s = pybamm.TertiaryBroadcast(N_s, [f"{domain} electrode"])
+            else:
+                D_eff = pybamm.SecondaryBroadcast(D_eff, [f"{domain} electrode"])
+                N_s = pybamm.SecondaryBroadcast(N_s, [f"{domain} electrode"])
+
         if self.size_distribution is True:
             # Size-dependent flux variables
             variables.update(
@@ -220,15 +227,10 @@ class FickianDiffusion(BaseParticle):
             )
             variables.update(self._get_standard_flux_distribution_variables(N_s))
             # Size-averaged flux variables
-            R = variables[f"{Phase_prefactor}{Domain} {phase_name}particle sizes [m]"]
+            D_eff = pybamm.size_average(D_eff)
+            R = pybamm.SpatialVariable("R", domains=D_eff.domains)
             f_a_dist = self.phase_param.f_a_dist(R)
-            D_eff = pybamm.Integral(f_a_dist * D_eff, R)
             N_s = pybamm.Integral(f_a_dist * N_s, R)
-
-        if self.x_average is True:
-            D_eff = pybamm.SecondaryBroadcast(D_eff, [f"{domain} electrode"])
-            N_s = pybamm.SecondaryBroadcast(N_s, [f"{domain} electrode"])
-
         variables.update(self._get_standard_diffusivity_variables(D_eff))
         variables.update(self._get_standard_flux_variables(N_s))
 
diff --git a/src/pybamm/models/submodels/particle/msmr_diffusion.py b/src/pybamm/models/submodels/particle/msmr_diffusion.py
index 8967116ce..4c73b29a4 100644
--- a/src/pybamm/models/submodels/particle/msmr_diffusion.py
+++ b/src/pybamm/models/submodels/particle/msmr_diffusion.py
@@ -240,6 +240,14 @@ class MSMRDiffusion(BaseParticle):
             }
         )
 
+        if self.x_average is True:
+            if self.size_distribution is True:
+                D_eff = pybamm.TertiaryBroadcast(D_eff, [f"{domain} electrode"])
+                N_s = pybamm.TertiaryBroadcast(N_s, [f"{domain} electrode"])
+            else:
+                D_eff = pybamm.SecondaryBroadcast(D_eff, [f"{domain} electrode"])
+                N_s = pybamm.SecondaryBroadcast(N_s, [f"{domain} electrode"])
+
         if self.size_distribution is True:
             # Size-dependent flux variables
             variables.update(
@@ -247,15 +255,10 @@ class MSMRDiffusion(BaseParticle):
             )
             variables.update(self._get_standard_flux_distribution_variables(N_s))
             # Size-averaged flux variables
-            R = variables[f"{Domain} {phase_name}particle sizes [m]"]
+            D_eff = pybamm.size_average(D_eff)
+            R = pybamm.SpatialVariable("R", domains=D_eff.domains)
             f_a_dist = self.phase_param.f_a_dist(R)
-            D_eff = pybamm.Integral(f_a_dist * D_eff, R)
             N_s = pybamm.Integral(f_a_dist * N_s, R)
-
-        if self.x_average is True:
-            D_eff = pybamm.SecondaryBroadcast(D_eff, [f"{domain} electrode"])
-            N_s = pybamm.SecondaryBroadcast(N_s, [f"{domain} electrode"])
-
         variables.update(self._get_standard_diffusivity_variables(D_eff))
         variables.update(self._get_standard_flux_variables(N_s))
 


Test patch:
diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py
index 7a728a744..a147104f4 100644
--- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py
+++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py
@@ -151,3 +151,16 @@ class TestDFNWithSizeDistribution:
         # compare
         np.testing.assert_array_almost_equal(neg_Li[0], neg_Li[1], decimal=12)
         np.testing.assert_array_almost_equal(pos_Li[0], pos_Li[1], decimal=12)
+
+    def test_basic_processing_nonlinear_diffusion(self):
+        options = {
+            "particle size": "distribution",
+        }
+        model = pybamm.lithium_ion.DFN(options)
+        # Ecker2015 has a nonlinear diffusion coefficient
+        parameter_values = pybamm.ParameterValues("Ecker2015")
+        parameter_values = pybamm.get_size_distribution_parameters(parameter_values)
+        modeltest = tests.StandardModelTest(
+            model, parameter_values=parameter_values, var_pts=self.var_pts
+        )
+        modeltest.test_all(skip_output_tests=True)
diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_mpm.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_mpm.py
index bbbf9b336..72a263c02 100644
--- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_mpm.py
+++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_mpm.py
@@ -122,3 +122,11 @@ class TestMPM:
         # compare
         np.testing.assert_array_almost_equal(neg_Li[0], neg_Li[1], decimal=13)
         np.testing.assert_array_almost_equal(pos_Li[0], pos_Li[1], decimal=13)
+
+    def test_basic_processing_nonlinear_diffusion(self):
+        model = pybamm.lithium_ion.MPM()
+        # Ecker2015 has a nonlinear diffusion coefficient
+        parameter_values = pybamm.ParameterValues("Ecker2015")
+        parameter_values = pybamm.get_size_distribution_parameters(parameter_values)
+        modeltest = tests.StandardModelTest(model, parameter_values=parameter_values)
+        modeltest.test_all(skip_output_tests=True)