Repository: PlasmaPy/PlasmaPy

Issue:
Pickle-able Particles (and other objects, for parallelization)
My draft Thomson scattering fitting function (#974) is currently restricted to using 1 CPU core because the underlying fit function (`thomson.spectral_density`) contains instances of the `Particle` object, which is not pickle-able because it includes a lambda function. If this issue was resolved, enabling parallelization of the fit function would be as simple as including a keyword to `scipy.optimize.differential_evolution`. 

Here is the error if I try to turn on parallelization now:

```
AttributeError: Can't pickle local object 'Particle.__init__.<locals>.<lambda>
```

and here is (I think) the offending line in `Particle.__init__`:

```
Line 408    -->    self._attributes = defaultdict(lambda: None)
``` 


Pickling a model that's then passed between multiple worker pools is, I gather, a pretty common way of parallelizing Python code, so I think it would be very desirable (maybe even a requirement) for all of these basic classes to be serializable. I know there are other serialization modules that DO support lambdas (eg. dill, cloudpickle), but they don't help since the issue is what's compatible with Scipy.

@namurphy  Is this something that can be easily changed?



Discussion:
This is really helpful, so thank you for bringing this up.  I started to address this in #1013, but removing `lambda` expressions in and of itself is not going to be enough to make `Particle` instances ~~pickleable~~ ~~picklable~~ able to be pickled. 

We might need to switch away from using a `defaultdict` for storing the attributes in `Particle` too.  We could do that if we set the following attributes to default to `None` manually in a regular `dict`, though it'd probably be better to refactor.

```Python
["half-life", "element", "isotope", "mass", "ion", "atomic number", "particle", "spin", "baryon number", "charge", "mass number",  "element name", "periodic table", "generation"]
```
I'm adding some (mostly failing) tests in #1014.  `DimensionlessParticle` and `CustomParticle` can be pickled, but `Particle` and ionization state stuff (which depends on `Particle`) cannot.
@namurphy I've realized that a similar problem will arise if we intend to use `numba` to pre-compile any of these functions, since in that case all of the function inputs need to be `numpy` objects. 

I was thinking that one way of solving this problem would be to split functions we want to optimize into a private optimized numba function and a public high-level wrapper that handles the units and translates objects like Particle. Here's a contrived example: 

```python
# A speedy, numba-ified subfunction whose inputs are all np types
# with assumed units.
@njit
def _some_function(x, Z, loc=3):
    return np.exp(- ((x-loc)/Z)**2)

# A slower wrapper public wrapper function that handles the nice Particle
# interface and astropy units
def some_function(x, loc =3*u.m, particle=Particle('He+')):
    Z = particle.atomic_number
    x = x.si.value
    loc = loc.si.value
    return _some_function(x, Z, loc=loc)
```

`_some_function` is compiled by numba, and is pickle-able, so we can use it internally in optimized code, while `some_function` still offers the nice user experience without any code being duplicated between them. We could split existing functions apart like this as needed: it wouldn't have to be done all at once. And it wouldn't change the public API at all.

Huh...that is an interesting idea!  This is potentially a strong use case of the `@from_particle` decorator proposed in #918.  

One thing to be careful about is that this would likely increase complexity and decrease readability, so I would tentatively reserve this approach for situations where optimization becomes necessary.  NumPy does pretty well with array operations but less well with iterative stuff, so my suspicion is that this could make a big difference for functions with loops in them but less of a difference for typical formulary functions.

One possibility would be to consolidate complexity in a decorator.  We could potentially create a new immutable `LeanParticle` object (possibly a `namedtuple`?) that contains only the mass, charge, integer charge, and particle symbol with no methods (and possibly no units).  If that gets created and passed through the `@particle_input` decorator (which might be given a `lean_particle=True` flag), then we could potentially use Numba on the inner stuff while preserving readability.  That approach would let us keep using the `Particle` API on the inside of functions too, which is an implementation characteristic I'd like to preserve.

I'm going to be working on more `plasmapy.particles` stuff in the near future, so I'd be happy to talk more.  I'm going to add this to my list of things to think about, in addition to some other important points that @rocco8773 has brought up with me recently too.  Thank you for bringing this up!
Moving complexity to the particle input decorator is an okay idea, but having it cover all possible use cases... Suppose we somehow (???) added numba capability there, and then realized that, I don't know, Yahoo made a library called Hax that offered an alternate and more efficient way of speeding up computation. In that case, we either keep pumping complexity into the decorator until it looks like the engineering version of W7X, or we keep using numba because of sunk costs.

The point of that extended metaphor is that I think Peter's approach is ultimately simpler, easier to maintain and generalizes to more use cases.
We would need to keep `@particle_input` separate from JIT decorators.  The responsibility of `@particle_input` would be to transform particle-like representations into the appropriate `Particle`, `CustomParticle`, or possibly something like a `LeanParticle` object, so that would be separate from JIT.  Sometimes we'll want to do one but not the other (like for functions that require one of the `Particle` methods rather than attributes, or computationally intensive methods that do not involve particles).

With respect to the problem of switching JIT decorators — one possibility that was recommended in either the `Clean Code` or `Clean Architecture` book would be to have a thin interface wrapper, perhaps something like: 
```Python
plasmapy._interfaces.jit = numba.jit
```
In each module we would do
```Python
from plasmapy._interfaces import jit

@jit
def reticulate_splines():
    ...
```
That way we would be able to change the JIT decorator throughout the whole package in a single place.  If we use keyword arguments in the decorator, we could use the bridge design pattern that would translate whatever JIT operator we would use into having the `numba.jit` API.  We would be tied to the Numba API, but not to Numba.  And if Numba is not installed, then we could fall back on an empty decorator that allows the Numba API.

With respect to `@particle_input`...what we'd need to do is ensure that the particle object that gets passed through is in a form that is compatible with a generic JIT operator when needed, which we would need to be careful about.
> With respect to @particle_input...what we'd need to do is ensure that the particle object that gets passed through is in a form that is compatible with a generic JIT operator when needed, which we would need to be careful about.

That last one sounds like something we could prototype and tinker with. But still, I can't really yet see a solution that doesn't condense to "take the necessary floats out of the particle object and pass them on".
Agreed — that is exactly what we would need it to do.  My perspective is that we could use a `namedtuple` that holds the relevant floats that were extracted from the particle object.  This would preserve the `Particle` API, while also putting the data in a form that should work with a generic JIT decorator.  For this I am making the assumption that a generic JIT decorator would be able to handle `namedtuple` objects, which seems probable but is something we'd need to verify before proceeding.
>  Named tuple classes, as returned by collections.namedtuple(), are supported in the same way regular tuples are supported. Attribute access and named parameters in the constructor are also supported.
>
> Creating a named tuple class inside Numba code is not supported; the class must be created at the global level.

from https://numba.pydata.org/numba-doc/dev/reference/pysupported.html#collections.

Dataclasses would be another way of doing that, but https://github.com/numba/numba/issues/4037 there's no support just yet.
> Dataclasses would be another way of doing that, but [numba/numba#4037](https://github.com/numba/numba/issues/4037) there's no support just yet.

Oh, that's a good point!  I had forgotten about dataclasses.  It appears I'm still stuck in Python 3.6+ thinking.  Dataclasses would probably be better in terms of allowing documentation for the different attributes, but a `namedtuple` could be a good starting point until Numba gets support for dataclasses.  

Along with #918 and #1079...the `@particle_input` refactoring in #1057 is going to need a lot of rethinking.  Oh well...getting things wrong is the first step towards getting things right!

Resolving patch:
diff --git a/plasmapy/particles/elements.py b/plasmapy/particles/elements.py
--- a/plasmapy/particles/elements.py
+++ b/plasmapy/particles/elements.py
@@ -14,9 +14,15 @@
 import json
 import pkgutil
 
-_PeriodicTable = collections.namedtuple(
-    "periodic_table", ["group", "category", "block", "period"]
-)
+from dataclasses import dataclass
+
+
+@dataclass
+class _PeriodicTable:
+    group: int
+    period: int
+    block: str
+    category: str
 
 
 def _element_obj_hook(obj):


Test patch:
diff --git a/plasmapy/particles/tests/test_pickling.py b/plasmapy/particles/tests/test_pickling.py
--- a/plasmapy/particles/tests/test_pickling.py
+++ b/plasmapy/particles/tests/test_pickling.py
@@ -19,17 +19,15 @@ class TestPickling:
     Test that different objects in `plasmapy.particles` can be pickled.
     """
 
-    xfail = pytest.mark.xfail(reason="see issue #1011")
-
     @pytest.mark.parametrize(
         "instance",
         [
             CustomParticle(mass=1 * u.kg, charge=1 * u.C),
             DimensionlessParticle(mass=5, charge=5),
-            pytest.param(Particle("p+"), marks=xfail),
-            pytest.param(IonicLevel("p+", 0.1, 1e9 * u.m ** -3), marks=xfail),
-            pytest.param(IonizationState("H", [0.5, 0.5]), marks=xfail),
-            pytest.param(IonizationStateCollection({"H": [0.5, 0.5]}), marks=xfail),
+            Particle("p+"),
+            IonicLevel("p+", 0.1, 1e9 * u.m ** -3),
+            IonizationState("H", [0.5, 0.5]),
+            IonizationStateCollection({"H": [0.5, 0.5]}),
         ],
     )
     def test_pickling(self, instance, tmp_path):