Extract kernel to a separate object using delegation

pyqg/__init__.py

@@ -5,6 +5,7 @@
 from .layered_model import LayeredModel
 from .particles import LagrangianParticleArray2D, GriddedLagrangianParticleArray2D
 from .parameterizations import *
+from .delegate import delegate
 
 try:
     from importlib.metadata import version, PackageNotFoundError

pyqg/bt_model.py

@@ -43,15 +43,13 @@ def __init__(self, beta=0.,  rd=0., H=1., U=0., **kwargs):
         self.Hi = np.array(H)[np.newaxis,...]
         self.U = U
 
-        self.nz = 1
-
         # deformation wavenumber
         if rd:
             self.kd2 = rd**-2
         else:
             self.kd2 = 0.
 
-        super().__init__(**kwargs)
+        super().__init__(nz=1, **kwargs)
 
         # initial conditions: (PV anomalies)
         self.set_q(1e-3*np.random.rand(1,self.ny,self.nx))

pyqg/delegate.py

@@ -0,0 +1,57 @@
+# Adapted from https://github.com/dscottboggs/python-delegate
+
+def delegate(*args, **named_args):
+    _dest = named_args.get('to')
+    if _dest is None:
+        raise ValueError(
+            "the named argument 'to' is required on the delegate function")
+
+    _prefix = named_args.get('prefix', '')
+
+    def wraps(cls, *wrapped_args, **wrapped_opts):
+        """Wrap the target class up in something that modifies."""
+        class Wrapped(cls):
+
+            _delegates = [(a, _dest, _prefix) for a in args] + getattr(cls, '_delegates', [])
+
+            def __getattr__(self, name):
+                """
+                Return the selected name from the destination if the name is one
+                of those selected. Since this is only called when `name` is not
+                found in `self.__dict__`, this method should always throw an
+                error when `name` is not one of the selected args to be
+                delegated.
+                """
+                for attr, dest, prefix in self._delegates:
+                    if name == prefix + attr:
+                        return getattr(self.__dict__[dest], name[len(prefix):])
+
+                raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'")
+
+            def __setattr__(self, name, value):
+                """
+                If this name is one of those selected, set it on the destination
+                property. Otherwise, set it on self.
+                """
+                for attr, dest, prefix in self._delegates:
+                    if name == prefix + attr:
+                        setattr(getattr(self, dest), name[len(prefix):], value)
+                        return
+                self.__dict__[name] = value
+
+            def __delattr__(self, name):
+                """Delete name from `dest` or `self`"""
+                for attr, dest, prefix in self._delegates:
+                    if name == prefix + attr:
+                        delattr(getattr(self, dest), name[len(prefix):])
+                        return
+
+                del self.__dict__[name]
+
+        Wrapped.__doc__ = cls.__doc__ or \
+            f"{cls.__class__} wrapped to delegate {args} to its {_dest} property"
+        Wrapped.__repr__ = cls.__repr__
+        Wrapped.__str__ = cls.__str__
+        Wrapped.__name__ = cls.__name__
+        return Wrapped
+    return wraps

pyqg/kernels/.compile_time_use_pyfftw.pxi

@@ -0,0 +1 @@
+DEF PYQG_USE_PYFFTW = 1

pyqg/kernels/__init__.py

@@ -0,0 +1 @@
+from .cython_fftw_kernel import CythonFFTWKernel