Skip to content

Commit 9f10f5e

Browse files
erin-atkinsonglwagner
erin-atkinson
and
glwagner
authored
Allow initialization of empty FieldDatasets similar to methods for FieldTimeSeries and define set! behaviour for FieldDataset (#4866)
* Add FieldDataset constructor for initialising sets of empty FieldTimeSeries * Add default behaviour for FieldDataset * Correct keyword argument ordering for FieldDataset * Remove unused methods * Default behaviour for individual fields * Remove breaking length check * `FieldDataset` constructor that uses information from existing fields * Add `set!` method for `FieldDataset` * Add `set!` test for `FieldDataset` * Fix use of `indices` and `boundary_conditions` * Add warning for unsupported metadata output * Add support for saving metadata with `FieldDataset` * Fix return convention for FieldDataset constructor Co-authored-by: Gregory L. Wagner <[email protected]> * Swap positional arguments to agree with `FieldTimeSeries` order Co-authored-by: Gregory L. Wagner <[email protected]> * Swap positional arguments to agree with `FieldTimeSeries` order Co-authored-by: Gregory L. Wagner <[email protected]> * Clean documentation for contructors * Add `set!` test for `InMemory` `FieldDataset` * Fix keyword argument whitespace Co-authored-by: Gregory L. Wagner <[email protected]> * Fix keyword argument whitespace Co-authored-by: Gregory L. Wagner <[email protected]> --------- Co-authored-by: Gregory L. Wagner <[email protected]> Co-authored-by: Gregory L. Wagner <[email protected]>
1 parent dd56694 commit 9f10f5e

File tree

2 files changed

+188
-2
lines changed

2 files changed

+188
-2
lines changed

‎src/OutputReaders/field_dataset.jl‎

Lines changed: 142 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
1-
struct FieldDataset{F, M, P, KW}
1+
using Oceananigans.Fields: instantiated_location, indices, boundary_conditions
2+
import Oceananigans.Fields: set!
3+
4+
struct FieldDataset{F, B, M, P, KW}
25
fields :: F
6+
backend :: B
37
metadata :: M
48
filepath :: P
59
reader_kw :: KW
@@ -52,7 +56,7 @@ function FieldDataset(filepath;
5256

5357
close(file)
5458

55-
return FieldDataset(ds, metadata, abspath(filepath), reader_kw)
59+
return FieldDataset(ds, backend, metadata, abspath(filepath), reader_kw)
5660
end
5761

5862
Base.getindex(fds::FieldDataset, inds...) = Base.getindex(fds.fields, inds...)
@@ -80,3 +84,139 @@ function Base.show(io::IO, fds::FieldDataset)
8084

8185
return print(io, s)
8286
end
87+
88+
"""
89+
FieldDataset(grid, times, fields;
90+
backend = OnDisk(),
91+
path = nothing,
92+
location = NamedTuple(),
93+
indices = NamedTuple(),
94+
boundary_conditions = NamedTuple(),
95+
metadata = Dict(),
96+
reader_kw = NamedTuple())
97+
98+
Returns a `FieldDataset` containing a new `FieldTimeSeries` for each key in `fields`
99+
on `grid` at `times`.
100+
101+
Keyword arguments
102+
=================
103+
- `backend`: backend, `InMemory(indices=Colon())` or `OnDisk()`
104+
105+
- `path`: path to data for `backend = OnDisk()`
106+
107+
- `location`: `NamedTuple` of location specifications, defaults to
108+
(Center, Center, Center) for each field`
109+
110+
- `indices`: `NamedTuple` of spatial indices, defaults to (:, :, :) for each field
111+
112+
- `boundary_conditions`: `NamedTuple` of boundary conditions for each field
113+
114+
- `metadata`: `Dict` containing metadata entries
115+
"""
116+
function FieldDataset(grid, times, fields::NTuple{N, Symbol};
117+
backend=OnDisk(),
118+
path=nothing,
119+
location=NamedTuple(),
120+
indices=NamedTuple(),
121+
boundary_conditions=NamedTuple(),
122+
metadata=Dict(),
123+
reader_kw=NamedTuple(),
124+
) where {N}
125+
126+
field_names = map(String, fields)
127+
128+
# Default behaviour
129+
indices = merge(
130+
NamedTuple(field=>(:, :, :) for field in fields),
131+
indices
132+
)
133+
location = merge(
134+
NamedTuple(field=>(Center(), Center(), Center()) for field in fields),
135+
location
136+
)
137+
boundary_conditions = merge(
138+
NamedTuple(field=>UnspecifiedBoundaryConditions() for field in fields),
139+
boundary_conditions
140+
)
141+
142+
# Create the FieldTimeSeries
143+
ftss = map(fields, field_names) do field, name
144+
inds = indices[field]
145+
loc = location[field]
146+
bcs = boundary_conditions[field]
147+
148+
FieldTimeSeries(loc, grid, times;
149+
indices=inds,
150+
backend,
151+
path,
152+
name,
153+
reader_kw,
154+
boundary_conditions=bcs
155+
)
156+
end
157+
158+
ds = Dict{String, FieldTimeSeries}(
159+
name => fts
160+
for (name, fts) in zip(field_names, ftss)
161+
)
162+
163+
return FieldDataset(ds, backend, metadata, path, reader_kw)
164+
end
165+
166+
167+
"""
168+
FieldDataset(times, fields;
169+
backend = OnDisk(),
170+
path = nothing,
171+
metadata = Dict(),
172+
reader_kw = NamedTuple())
173+
174+
Returns a `FieldDataset` containing a new `FieldTimeSeries` for each field
175+
in the `NamedTuple``fields` at `times`. Locations, indices and boundary
176+
conditions are extracted from `fields``
177+
178+
Keyword arguments
179+
=================
180+
- `backend`: backend, `InMemory(indices=Colon())` or `OnDisk()`
181+
182+
- `path`: path to data for `backend = OnDisk()`
183+
184+
- `metadata`: `Dict` containing metadata entries
185+
"""
186+
function FieldDataset(times, fields;
187+
fds_kw...
188+
)
189+
190+
grid = fields[1].grid
191+
any([field.grid != grid for field in fields]) && throw(ArgumentError("All fields must be defined on the same grid"))
192+
193+
loc = map(instantiated_location, fields)
194+
inds = map(Fields.indices, fields)
195+
bcs = map(Fields.boundary_conditions, fields)
196+
197+
return FieldDataset(grid, times, keys(fields);
198+
location=loc,
199+
indices=inds,
200+
boundary_conditions=bcs,
201+
fds_kw...
202+
)
203+
end
204+
205+
# Setting a FieldDataset iterates over contained FieldTimeSeries
206+
function set!(fds::FieldDataset, args...; fields...)
207+
for (k, v) in pairs(fields)
208+
set!(fds[k], v, args...)
209+
end
210+
end
211+
212+
# Write metadata if possible for OnDisk FieldDataset
213+
function set!(fds::FieldDataset{F, B, M, P, KW}, args...; fields...) where {F, B<:OnDisk, M, P, KW}
214+
jldopen(fds.filepath, "a+") do file
215+
for (k, v) in pairs(fds.metadata)
216+
maybe_write_property!(file, "metadata/$k", v)
217+
end
218+
end
219+
for (k, v) in pairs(fields)
220+
set!(fds[k], v, args...)
221+
end
222+
end

‎test/test_output_readers.jl‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,52 @@ end
456456
@test t[1, 1, 1] == 3.8
457457
end
458458

459+
@testset "Outputwriting with set!(FieldDataset{OnDisk})" begin
460+
@info " Testing set!(FieldDataset{OnDisk})..."
461+
462+
grid = RectilinearGrid(size = (1, 1, 1), extent = (1, 1, 1))
463+
464+
a = CenterField(grid)
465+
b = Field{Face, Center, Center}(grid)
466+
467+
metadata = Dict("i"=>12, "j"=>"jay")
468+
469+
filepath = "testfile.jld2"
470+
f = FieldDataset(1:10, (; a, b); backend = OnDisk(), path = filepath, metadata)
471+
472+
for i in 1:10
473+
set!(a, i)
474+
set!(b, 2i)
475+
set!(f, i; a, b)
476+
end
477+
478+
g = FieldDataset(filepath)
479+
480+
@test location(g.a) == (Center, Center, Center)
481+
@test location(g.b) == (Face, Center, Center)
482+
@test g.a.grid == a.grid
483+
@test g.b.grid == b.grid
484+
485+
@test g.a[1, 1, 1, 1] == 1
486+
@test g.a[1, 1, 1, 10] == 10
487+
@test g.a[1, 1, 1, Time(1.6)] == 1.6
488+
489+
@test g.b[1, 1, 1, 1] == 2
490+
@test g.b[1, 1, 1, 10] == 20
491+
@test g.b[1, 1, 1, Time(5.1)] == 10.2
492+
493+
@test g.metadata["i"] == 12
494+
@test g.metadata["j"] == "jay"
495+
496+
t = g.a[Time(3.8)]
497+
498+
@test t[1, 1, 1] == 3.8
499+
500+
set!(g, 2; a=-1, b=-2)
501+
@test g.a[1, 1, 1, 2] == -1
502+
@test g.b[1, 1, 1, 2] == -2
503+
end
504+
459505
@testset "Test chunked abstraction" begin
460506
@info " Testing Chunked abstraction..."
461507
filepath = "testfile.jld2"

0 commit comments

Comments
 (0)