Repository: zwicker-group/py-pde

Issue:
Add option for transformer function to `StorageTracker`
`StorageTracker` should not only be able to store the actual fields of the simulations, but also some derived fields. The idea would be to add an argument to the StoragerTracker which could take a function that computes a field or field collection based on the current state and time. The default would be to simply forward the field.

The idea was raised in #433 

Discussion:

Patch:
diff --git a/pde/storage/base.py b/pde/storage/base.py
index 10ec98c..d5d12aa 100644
--- a/pde/storage/base.py
+++ b/pde/storage/base.py
@@ -266,18 +266,30 @@ class StorageBase(metaclass=ABCMeta):
 
     @fill_in_docstring
     def tracker(
-        self, interval: Union[int, float, InterruptsBase] = 1
+        self,
+        interval: Union[int, float, InterruptsBase] = 1,
+        *,
+        transformation: Optional[Callable[[FieldBase, float], FieldBase]] = None,
     ) -> "StorageTracker":
         """create object that can be used as a tracker to fill this storage
 
         Args:
             interval:
                 {ARG_TRACKER_INTERVAL}
+            transformation (callable, optional):
+                A function that transforms the current state into a new field or field
+                collection, which is then stored. This allows to store derived
+                quantities of the field during calculations. The argument needs to be a
+                callable function taking 1 or 2 arguments. The first argument always is
+                the current field, while the optional second argument is the associated
+                time.
 
         Returns:
             :class:`StorageTracker`: The tracker that fills the current storage
         """
-        return StorageTracker(storage=self, interval=interval)
+        return StorageTracker(
+            storage=self, interval=interval, transformation=transformation
+        )
 
     def start_writing(self, field: FieldBase, info: Optional[InfoDict] = None) -> None:
         """initialize the storage for writing data
@@ -494,16 +506,41 @@ class StorageTracker(TrackerBase):
     """
 
     @fill_in_docstring
-    def __init__(self, storage, interval: IntervalData = 1):
+    def __init__(
+        self,
+        storage,
+        interval: IntervalData = 1,
+        *,
+        transformation: Optional[Callable[[FieldBase, float], FieldBase]] = None,
+    ):
         """
         Args:
             storage (:class:`~pde.storage.base.StorageBase`):
                 Storage instance to which the data is written
             interval:
                 {ARG_TRACKER_INTERVAL}
+            transformation (callable, optional):
+                A function that transforms the current state into a new field or field
+                collection, which is then stored. This allows to store derived
+                quantities of the field during calculations. The argument needs to be a
+                callable function taking 1 or 2 arguments. The first argument always is
+                the current field, while the optional second argument is the associated
+                time.
         """
         super().__init__(interval=interval)
         self.storage = storage
+        if transformation is not None and not callable(transformation):
+            raise TypeError("`transformation` must be callable")
+        self.transformation = transformation
+
+    def _transform(self, field: FieldBase, t: float) -> FieldBase:
+        """transforms the field according to the defined transformation"""
+        if self.transformation is None:
+            return field
+        elif self.transformation.__code__.co_argcount == 1:
+            return self.transformation(field)  # type: ignore
+        else:
+            return self.transformation(field, t)
 
     def initialize(self, field: FieldBase, info: Optional[InfoDict] = None) -> float:
         """
@@ -517,7 +554,7 @@ class StorageTracker(TrackerBase):
             float: The first time the tracker needs to handle data
         """
         result = super().initialize(field, info)
-        self.storage.start_writing(field, info)
+        self.storage.start_writing(self._transform(field, 0), info)
         return result
 
     def handle(self, field: FieldBase, t: float) -> None:
@@ -528,7 +565,7 @@ class StorageTracker(TrackerBase):
                 The current state of the simulation
             t (float): The associated time
         """
-        self.storage.append(field, time=t)
+        self.storage.append(self._transform(field, t), time=t)
 
     def finalize(self, info: Optional[InfoDict] = None) -> None:
         """finalize the tracker, supplying additional information
diff --git a/pde/storage/file.py b/pde/storage/file.py
index 8e871be..0fcc798 100644
--- a/pde/storage/file.py
+++ b/pde/storage/file.py
@@ -9,7 +9,7 @@ from __future__ import annotations
 import json
 import logging
 from pathlib import Path
-from typing import Any, Optional, Tuple  # @UnusedImport
+from typing import Any, Optional, Tuple
 
 import numpy as np
 from numpy.typing import DTypeLike
@@ -26,8 +26,8 @@ class FileStorage(StorageBase):
     def __init__(
         self,
         filename: str,
-        info: Optional[InfoDict] = None,
         *,
+        info: Optional[InfoDict] = None,
         write_mode: str = "truncate_once",
         max_length: Optional[int] = None,
         compression: bool = True,
diff --git a/pde/storage/memory.py b/pde/storage/memory.py
index 51aa264..aa4e9bc 100644
--- a/pde/storage/memory.py
+++ b/pde/storage/memory.py
@@ -23,8 +23,9 @@ class MemoryStorage(StorageBase):
         self,
         times: Optional[Sequence[float]] = None,
         data: Optional[List[np.ndarray]] = None,
-        field_obj: Optional[FieldBase] = None,
+        *,
         info: Optional[InfoDict] = None,
+        field_obj: Optional[FieldBase] = None,
         write_mode: str = "truncate_once",
     ):
         """
diff --git a/pde/tools/numba.py b/pde/tools/numba.py
index af8c50b..39158d0 100644
--- a/pde/tools/numba.py
+++ b/pde/tools/numba.py
@@ -106,7 +106,7 @@ def numba_environment() -> Dict[str, Any]:
         threading_layer = nb.threading_layer()
     except ValueError:
         # threading layer was not initialized, so compile a mock function
-        @nb.jit("i8()", parallel=True)
+        @nb.njit("i8()", parallel=True)
         def f():
             s = 0
             for i in nb.prange(4):


Test patch:
diff --git a/tests/storage/test_generic_storages.py b/tests/storage/test_generic_storages.py
index a8da6d6..465659d 100644
--- a/tests/storage/test_generic_storages.py
+++ b/tests/storage/test_generic_storages.py
@@ -258,3 +258,45 @@ def test_storage_mpi(storage_factory):
     if mpi.is_main:
         assert res.integral == pytest.approx(field.integral)
         assert len(storage) == 11
+
+
+@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
+def test_storing_transformation_collection(storage_factory):
+    """test transformation yielding field collections in storage classes"""
+    grid = UnitGrid([8])
+    field = ScalarField.random_normal(grid).smooth(1)
+
+    def trans1(field, t):
+        return FieldCollection([field, 2 * field + t])
+
+    storage = storage_factory()
+    eq = DiffusionPDE()
+    trackers = [storage.tracker(0.01, transformation=trans1)]
+    eq.solve(
+        field,
+        t_range=0.1,
+        dt=0.001,
+        backend="numpy",
+        tracker=trackers,
+    )
+
+    assert storage.has_collection
+    for t, sol in storage.items():
+        a, a2 = sol
+        np.testing.assert_allclose(a2.data, 2 * a.data + t)
+
+
+@pytest.mark.parametrize("storage_class", STORAGE_CLASSES)
+def test_storing_transformation_scalar(storage_factory):
+    """test transformations yielding scalar fields in storage classes"""
+    grid = UnitGrid([8])
+    field = ScalarField.random_normal(grid).smooth(1)
+
+    storage = storage_factory()
+    eq = DiffusionPDE(diffusivity=0)
+    trackers = [storage.tracker(0.01, transformation=lambda f: f**2)]
+    eq.solve(field, t_range=0.1, dt=0.001, backend="numpy", tracker=trackers)
+
+    assert not storage.has_collection
+    for sol in storage:
+        np.testing.assert_allclose(sol.data, field.data**2)