Support FieldTimeSeries targets in Relaxation

[ewquon]

Summary

• Relaxation(rate=…, target=fts) now accepts a FieldTimeSeries as the relaxation target.
• At materialize time, the FTS is wrapped in an internal FieldTimeSeriesTarget{L,F} carrying the simulation-side location (instantiated_location(field)) and the integer index of the forced field in model_fields. The Relaxation struct itself is unchanged (still Relaxation{R,M,T}); the target field's type just becomes FieldTimeSeriesTarget after materialization.
• A new kernel callable on Relaxation{R,M,<:FieldTimeSeriesTarget} reads ϕ = model_fields[index][i,j,k] and obtains ϕᵣ via interpolate(X, Time(clock.time), fts, instantiated_location(fts), fts.grid) (spatial + temporal interpolation), so the FTS may live on a different grid as long as its extent brackets the simulation grid in x, y, and z.
• Mismatched extents throw ArgumentError at materialize time naming the offending axis and both intervals.

The existing Relaxation(target=callable) / Relaxation(target=number) path is untouched — multiple dispatch routes those cases through the original ContinuousForcing wrapping.

Why

This is the first slice of a small series that generalizes Relaxation for real-data use cases — driving an LES or regional run with prescribed large-scale state from a host model (ERA5, GFS, a parent simulation) through grid nudging and/or open boundaries with a fringe region. Scope of this PR is intentionally narrow: 3D mode only, FTS targets only. "Profile mode" (mode=:profile + column reference, see NumericalEarth/Breeze.jl#562) and other target types are deferred to follow-up PRs.

Design notes

• Mirrors the two existing patterns it composes: Forcing(fts::FlavorOfFTS) materializes into a discrete forcing (forcing.jl:184-186); Relaxation already mutates the user's recipe into a kernel-ready form at materialize_forcing (relaxation.jl:75-78). This PR keeps the Relaxation identity through materialization (model.forcing.c isa Relaxation stays true), only swapping the target field's type.
• Field index stored as runtime Int, matching the precedent in ContinuousForcing.field_dependencies_indices (user_function_arguments.jl:1-19) rather than encoding as a Val{N} type param.
• Existing Forcing(fts) doesn't enforce location matching (flagged latent constraint); this PR sidesteps that by using interpolate, so location/grid alignment between FTS and forced field is no longer required.

Test plan

• CPU unit tests pass locally (6/6: materialization wiring + analytical convergence + extent validation)
• Existing Relaxation smoke tests (GaussianMask, PiecewiseLinearMask) still pass
• Doctest output verified to match
• CI green (CPU + GPU)
• Reactant trace smoke test (deferred to follow-up unless CI catches a regression)

Out of scope (follow-ups)

• Interpolation performance optimization by caching FractionalIndices arrays in FieldTimeSeriesTarget at materialize time.
• Profile mode (mode=:profile, column-averaged relaxation toward 1×1×Nz reference). Requires the new update_forcing! hook in each model's update_state!.
• Plain Field (time-invariant) targets.
• column_field_time_series utility for extracting a column reference from a 3D reanalysis FTS.

[glwagner]

just a few minor aesthetic points

[glwagner]

why is this draft?

[simone-silvestri]

Do we need a new type? I think we can achieve the same thing in a simpler and more general way by extending Relaxation to wrap Fields and FieldTimeSeries in a DiscreteForcings rather than ContinuousForcing (maybe a type parameter on Relaxation to decide wether it's discrete or continuous). Then we can create simple function that grab the correct Time in the field time series. This means that we will also support Field targets without issues. Something like this might be enough

DiscreteRelaxation = Relaxation{<:Any, <:Any, <:Union{AbstractArray, FlavorOfFTS}}

@inline get_target(i, j, k, grid, a::AbstactArray, clock) = @inbounds a[i, j, k]
@inline get_target(i, j, k, grid, a::FlavorOfFTS, clock) = @inbounds a[i, j, k, Time(clock)]

function (r::DiscreteRelaxation)(i, j, k, grid, clock, fields, params) 
      target = get_target(i, j, k, grid, r.target, clock)
      φ = @inbounds fields[p.field_name][i, j, k]
      X = node(i, j, k, grid, location(p.fts)...)
      return r.rate * r.mask(X...) * (target - φ)
end

function materialize_forcing(forcing::DiscreteRelaxation, field, field_name, model_field_names)
      discrete_relaxation = DiscreteForcing(forcing; parameters=(; field_name))
      return materialize_forcing(discrete_relaxation, field, field_name, model_field_names)
end

but probably better to think a bit more about it such that is a bit more general.

[simone-silvestri]

but maybe to do this we need to pass the location as parameters to the DiscreteForcing as this is thrown away in case of a Field target when adapting to GPU

[ewquon]

why is this draft?

@glwagner I was looking into optimizing the interpolation by caching the fractional indices but I think we should keep it simple.

[glwagner]

Do we need a new type? I think we can achieve the same thing in a simpler and more general way by extending Relaxation to wrap Fields and FieldTimeSeries in a DiscreteForcings rather than ContinuousForcing (maybe a type parameter on Relaxation to decide wether it's discrete or continuous). Then we can create simple function that grab the correct Time in the field time series. This means that we will also support Field targets without issues. Something like this might be enough

DiscreteRelaxation = Relaxation{<:Any, <:Any, <:Union{AbstractArray, FlavorOfFTS}}

@inline get_target(i, j, k, grid, a::AbstactArray, clock) = @inbounds a[i, j, k]
@inline get_target(i, j, k, grid, a::FlavorOfFTS, clock) = @inbounds a[i, j, k, Time(clock)]

function (r::DiscreteRelaxation)(i, j, k, grid, clock, fields, params) 
      target = get_target(i, j, k, grid, r.target, clock)
      φ = @inbounds fields[p.field_name][i, j, k]
      X = node(i, j, k, grid, location(p.fts)...)
      return r.rate * r.mask(X...) * (target - φ)
end

function materialize_forcing(forcing::DiscreteRelaxation, field, field_name, model_field_names)
      discrete_relaxation = DiscreteForcing(forcing; parameters=(; field_name))
      return materialize_forcing(discrete_relaxation, field, field_name, model_field_names)
end

but probably better to think a bit more about it such that is a bit more general.

The new target type is needed to support interpolation of the FieldTimeSeries. This is needed to support relaxation to FieldTimeSeries on coarse grids (a typical situation that we encounter for open boundary conditions). The other issue is that we need to encode the location of the target. For example, ERA5 velocities are located at cell centers, but Oceananigans and Breeze velocities are on the C-grid.

[glwagner]

I do think we could use DiscreteForcing with parameters rather than a new target type though, as an alternative implementation with slightly more code re-use. But @simone-silvestri's design will not work out of the box:

1. We cannot pass field_name::Symbol in parameters, because Symbol are not isbits. However, we can compute the field index and pass that to parameters instead.
2. We additionally need the location / instantiated location in parameters, in order to interpolate.

Switching to this design does preclude the future possibility of an optimization that embeds the field index as a type parameter (so the compiler can infer fields[field_index]). I don't know if that would matter.

The DiscreteForcing design might also need to be undone later, if we want to support transformations on the source field (ie if we want to relax the domain average field, rather than the 3D field directly). But, that could be part of a future PR if desired.

[simone-silvestri]

We cannot pass field_name::Symbol in parameters, because Symbol are not isbits. However, we can compute the field index and pass that to parameters instead.

I think this points to a flaw in the forcing signature: we should pass the forced field to the forcing function.
This might also allow to simplify the implementation of DatasetRestoring in NumericalEarth.

[ewquon]

@simone-silvestri I really appreciate your time in reviewing this and providing feedback. For the time being, I'd like to move forward with what's currently in the PR to provide some additional flexibility as use cases evolve—I'd rather not refactor to DiscreteForcing now and then undo it later. Happy to revisit and refactor down the road as necessary.

[simone-silvestri]

I would argue that this is not a general solution. The general solution would be fixing the technical dept that require us to go in hoops to pass the field name in the forcing. This solution layers on top of this flaw, so the best course of action would be to solve this issue.

What I propose would be a more substantial refactor but would prove a little more solid in the future. I understand, however if you need this feature immediately. We can refactor after this.

[simone-silvestri]

Note that passing the forced field in the forcing function signature is basically what ContinuousForcing is already doing (through a slightly more convoluted path) this is why we don't need all this extra boilerplate to have a simple relaxation in that case.

[glwagner]

I think we should refactor in a future PR to remove the ContinuousForcing dependency. Another thing we want to add is to allow transformation of field. This would motivate simply storing the field in the struct directly (rather than trying to extract it from model_fields.) The sum of those two changes would be a decent simplification for a future PR.

[glwagner]

sadly GPU tests are now broken...