From 6a377c8865607b0e7b73ef72c6a85edcaa57379e Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Tue, 17 Dec 2024 12:53:13 -0500 Subject: [PATCH 1/4] Make copies of direct-Field and duplicate tasks to avoid processing issues and transform noise in coeff outputs --- dedalus/core/evaluator.py | 29 +++++++++++++---------------- dedalus/core/operators.py | 4 +++- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/dedalus/core/evaluator.py b/dedalus/core/evaluator.py index 2a26dc12..f5000208 100644 --- a/dedalus/core/evaluator.py +++ b/dedalus/core/evaluator.py @@ -13,6 +13,7 @@ from .future import FutureField, FutureLockedField from .field import Field, LockedField +from .operators import Copy from ..tools.cache import CachedAttribute from ..tools.general import OrderedSet from ..tools.general import oscillate @@ -127,22 +128,19 @@ def evaluate_handlers(self, handlers, id=None, **kw): # Attempt evaluation tasks = self.attempt_tasks(tasks, id=id) - # # Transform all outputs to coefficient layout to dealias - ## D3 note: need to worry about this for redundent tasks? - # outputs = OrderedSet([t['out'] for h in handlers for t in h.tasks]) - # self.require_coeff_space(outputs) - - # # Copy redundant outputs so processing is independent - # outputs = set() - # for handler in handlers: - # for task in handler.tasks: - # if task['out'] in outputs: - # task['out'] = task['out'].copy() - # else: - # outputs.add(task['out']) + # Transform all outputs to coefficient layout to dealias outputs = OrderedSet([t['out'] for h in handlers for t in h.tasks if not isinstance(t['out'], LockedField)]) self.require_coeff_space(outputs) + # Copy redundant outputs so processing is independent + outputs = set() + for handler in handlers: + for task in handler.tasks: + if task['out'] in outputs: + task['out'] = task['out'].copy() + else: + outputs.add(task['out']) + # Process for handler in handlers: handler.process(**kw) @@ -285,10 +283,9 @@ def add_task(self, task, layout='g', name=None, scales=None): # Create operator if isinstance(task, str): op = FutureField.parse(task, self.vars, self.dist) + elif isinstance(task, Field): + op = Copy(task) else: - # op = FutureField.cast(task, self.domain) - # op = Cast(task) - # TODO: figure out if we need to copying here op = task # Check scales if isinstance(op, (LockedField, FutureLockedField)): diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 09e66ff6..9a9a993d 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -19,7 +19,7 @@ from .domain import Domain from . import coords -from .field import Operand, Field +from .field import Operand, Field, LockedField from .future import Future, FutureField, FutureLockedField from ..tools.array import reshape_vector, apply_matrix, add_sparse, axindex, axslice, perm_matrix, copyto, sparse_block_diag, interleave_matrices from ..tools.cache import CachedAttribute, CachedMethod @@ -1492,6 +1492,8 @@ class Copy(LinearOperator): def __init__(self, operand, out=None): super().__init__(operand, out=out) + if isinstance(operand, (LockedField, FutureLockedField)): + raise ValueError("Not yet implemented for locked fields.") # LinearOperator requirements self.operand = operand # FutureField requirements From 94c21e5f65284bf71e7a922b10222388ae1e1bdb Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Tue, 17 Dec 2024 15:56:25 -0500 Subject: [PATCH 2/4] Remove field locks in NLBVP linearization --- dedalus/core/field.py | 7 +++++++ dedalus/core/problems.py | 6 +++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/dedalus/core/field.py b/dedalus/core/field.py index c7359fb6..a2ac3489 100644 --- a/dedalus/core/field.py +++ b/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 + diff --git a/dedalus/core/problems.py b/dedalus/core/problems.py index 98daeea8..5da537f6 100644 --- a/dedalus/core/problems.py +++ b/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) From 5d7a1e9508a7fa8af9b712f24ef2a570fcfdd966 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 18 Dec 2024 14:44:04 -0500 Subject: [PATCH 3/4] Move hermitian enforcement to after timestep, so it isn't applied twice when reloading on hermitian project steps --- dedalus/core/solvers.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/dedalus/core/solvers.py b/dedalus/core/solvers.py index ddd27155..c63a5fbb 100644 --- a/dedalus/core/solvers.py +++ b/dedalus/core/solvers.py @@ -685,11 +685,6 @@ def step(self, dt): # Assert finite timestep if not np.isfinite(dt): raise ValueError("Invalid timestep") - # Enforce Hermitian symmetry for real variables - if self.enforce_real_cadence: - # Enforce for as many iterations as timestepper uses internally - if self.iteration % self.enforce_real_cadence < self.timestepper.steps: - self.enforce_hermitian_symmetry(self.state) # Record times wall_time = self.wall_time if self.iteration == self.initial_iteration: @@ -706,6 +701,11 @@ def step(self, dt): self.run_time_start = self.wall_time # Advance using timestepper self.timestepper.step(dt, wall_time) + # Enforce Hermitian symmetry for real variables + if self.enforce_real_cadence: + # Enforce for as many iterations as timestepper uses internally + if self.iteration % self.enforce_real_cadence < self.timestepper.steps: + self.enforce_hermitian_symmetry(self.state) # Update iteration self.iteration += 1 self.dt = dt From 14755154660c106db95136c5013d3ad44cc698b4 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Wed, 18 Dec 2024 14:44:58 -0500 Subject: [PATCH 4/4] Save weak references to all fields built on a distributor, to help track fields and memory usage --- dedalus/core/distributor.py | 3 +++ dedalus/core/field.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/dedalus/core/distributor.py b/dedalus/core/distributor.py index db63a8c7..c4cc766f 100644 --- a/dedalus/core/distributor.py +++ b/dedalus/core/distributor.py @@ -9,6 +9,7 @@ from collections import OrderedDict from math import prod import numbers +from weakref import WeakSet from .coords import CoordinateSystem, DirectProduct from ..tools.array import reshape_vector @@ -112,6 +113,8 @@ def __init__(self, coordsystems, comm=None, mesh=None, dtype=None): self.comm_coords = np.array(self.comm_cart.coords, dtype=int) # Build layout objects self._build_layouts() + # Keep set of weak field references + self.fields = WeakSet() @CachedAttribute def cs_by_axis(self): diff --git a/dedalus/core/field.py b/dedalus/core/field.py index a2ac3489..415edcf6 100644 --- a/dedalus/core/field.py +++ b/dedalus/core/field.py @@ -572,6 +572,8 @@ def __init__(self, dist, bases=None, name=None, tensorsig=None, dtype=None): self.layout = self.dist.get_layout_object('c') # Change scales to build buffer and data self.preset_scales((1,) * self.dist.dim) + # Add weak reference to distributor + dist.fields.add(self) def __getitem__(self, layout): """Return data viewed in specified layout."""