Remove field locks in NLBVP linearization

Author: kburns

dedalus/core/field.py

@@ -1022,3 +1022,10 @@ def lock_to_layouts(self, *layouts):
     def lock_axis_to_grid(self, axis):
         self.allowed_layouts = tuple(l for l in self.dist.layouts if l.grid_space[axis])
 
+    def unlock(self):
+        """Return regular Field object with same data and no layout locking."""
+        field = Field(self.dist, bases=self.domain.bases, name=self.name, tensorsig=self.tensorsig, dtype=self.dtype)
+        field.preset_scales(self.scales)
+        field[self.layout] = self.data
+        return field
+

dedalus/core/problems.py

@@ -3,7 +3,7 @@
 import numpy as np
 from collections import ChainMap
 
-from .field import Operand, Field
+from .field import Operand, Field, LockedField
 from . import arithmetic
 from . import operators
 from . import solvers
@@ -244,6 +244,10 @@ def _build_matrix_expressions(self, eqn):
         # Extract matrix expressions
         F = eqn['LHS'] - eqn['RHS']
         dF = F.frechet_differential(vars, perts)
+        # Remove any field locks
+        dF = dF.replace(operators.Lock, lambda x: x)
+        for field in dF.atoms(LockedField):
+            dF = dF.replace(field, field.unlock())
         # Reinitialize and prep NCCs
         dF = dF.reinitialize(ncc=True, ncc_vars=perts)
         dF.prep_nccs(vars=perts)