From a70ebc6ac14371b8fac109589b9e289bb2aae76a Mon Sep 17 00:00:00 2001
From: "Gregory L. Wagner" <wagner.greg@gmail.com>
Date: Sat, 14 Dec 2024 11:56:36 -0700
Subject: [PATCH] Add support for loading data generated prior to 0.95.0 with
 FieldTimeSeries (#3990)

* Try to manually build grids

* Add shortcut for building new grids in FieldTimeSeries
---
 src/OutputReaders/field_time_series.jl     | 65 +++++++++++++++++++---
 src/OutputReaders/set_field_time_series.jl |  1 +
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl
index 19a1bdb2c8..afe9439e38 100644
--- a/src/OutputReaders/field_time_series.jl
+++ b/src/OutputReaders/field_time_series.jl
@@ -13,7 +13,8 @@ using Oceananigans.Architectures
 using Oceananigans.Grids
 using Oceananigans.Fields
 
-using Oceananigans.Grids: topology, total_size, interior_parent_indices, parent_index_range
+using Oceananigans.Grids: topology, total_size, interior_parent_indices, parent_index_range, AbstractGrid
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
 
 using Oceananigans.Fields: interior_view_indices, index_binary_search,
                            indices_summary, boundary_conditions
@@ -386,6 +387,10 @@ function FieldTimeSeries(loc, grid, times=();
                                        times, path, name, time_indexing, reader_kw)
 end
 
+isreconstructed(grid::JLD2.ReconstructedStatic) = true
+isreconstructed(grid::AbstractGrid) = false
+isreconstructed(grid::ImmersedBoundaryGrid) = isreconstructed(grid.underlying_grid)
+
 """
     FieldTimeSeries{LX, LY, LZ}(grid::AbstractGrid [, times=()]; kwargs...)
 
@@ -466,8 +471,6 @@ function FieldTimeSeries(path::String, name::String;
         (:, :, :)
     end
 
-    isnothing(grid) && (grid = file["serialized/grid"])
-
     if isnothing(architecture) # determine architecture
         if isnothing(grid) # go to default
             architecture = CPU()
@@ -476,11 +479,54 @@ function FieldTimeSeries(path::String, name::String;
         end
     end
 
-    if boundary_conditions isa UnspecifiedBoundaryConditions
-        boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"]
-        boundary_conditions = on_architecture(architecture, boundary_conditions)
-    end
+    isnothing(grid) && (grid = file["serialized/grid"])
 
+    # If isreconstructed(grid), it probably means that the data was generated prior to
+    # Oceananigans version 0.95.0 (12/13/2024). In this case, the best we can do is to try to rebuild
+    # the grids manually using the non-serialized grid data. Here, we support RectilinearGrid
+    # and LatitudeLongitudeGrid (but not OrthogonalSphericalShellGrid) and we also assume
+    # GridFittedBottom if the grid is an ImmersedBoundaryGrid. If these assumptions can be relaxed
+    # in the future, they should.
+    if isreconstructed(grid)
+        isibg = grid isa ImmersedBoundaryGrid
+        test_grid = isibg ? grid.underlying_grid : grid
+        address = isibg ? "grid/underlying_grid" : "grid"
+        Nx = file["$address/Nx"]
+        Ny = file["$address/Ny"]
+        Nz = file["$address/Nz"]
+        Hx = file["$address/Hx"]
+        Hy = file["$address/Hy"]
+        Hz = file["$address/Hz"]
+        zᵃᵃᶠ = file["$address/zᵃᵃᶠ"]
+        z = file["$address/Δzᵃᵃᶠ"] isa Number ? (zᵃᵃᶠ[1], zᵃᵃᶠ[Nz+1]) : zᵃᵃᶠ[1:Nz+1]
+        topo = topology(grid)
+        size = (Nx, Ny, Nz)
+        halo = (Hx, Hy, Hz)
+
+        if :λᶜᵃᵃ ∈ propertynames(test_grid) # I guess its a LatitudeLongitudeGrid.
+            λᶠᵃᵃ = file["$address/λᶠᵃᵃ"]
+            φᵃᶠᵃ = file["$address/φᵃᶠᵃ"]
+            λ = file["$address/Δλᶠᵃᵃ"] isa Number ? (λᶠᵃᵃ[1], λᶠᵃᵃ[Nx+1]) : λᶠᵃᵃ[1:Nx+1]
+            φ = file["$address/Δφᵃᶠᵃ"] isa Number ? (φᵃᶠᵃ[1], φᵃᶠᵃ[Ny+1]) : φᵃᶠᵃ[1:Ny+1]
+            domain = (latitude=φ, longitude=λ, z=z)
+            underlying_grid = LatitudeLongitudeGrid(architecture; size, halo, topology=topo, domain...)
+        else
+            xᶠᵃᵃ = file["$address/xᶠᵃᵃ"]
+            yᵃᶠᵃ = file["$address/yᵃᶠᵃ"]
+            x = file["$address/Δxᶠᵃᵃ"] isa Number ? (xᶠᵃᵃ[1], xᶠᵃᵃ[Nx+1]) : xᶠᵃᵃ[1:Nx+1]
+            y = file["$address/Δyᵃᶠᵃ"] isa Number ? (yᵃᶠᵃ[1], yᵃᶠᵃ[Ny+1]) : yᵃᶠᵃ[1:Ny+1]
+            domain = (; x, y, z)
+            underlying_grid = RectilinearGrid(architecture; size, halo, topology=topo, domain...)
+        end
+
+        if isibg
+            bottom_height = file["grid/immersed_boundary/bottom_height"]
+            bottom_height = view(bottom_height, 1+Hx:Nx+Hx, 1+Hy:Ny+Hy, 1)
+            grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
+        else
+            grid = underlying_grid
+        end
+    end
 
     # This should be removed eventually... (4/5/2022)
     grid = try
@@ -529,6 +575,11 @@ function FieldTimeSeries(path::String, name::String;
         end
     end
 
+    if boundary_conditions isa UnspecifiedBoundaryConditions
+        boundary_conditions = file["timeseries/$name/serialized/boundary_conditions"]
+        boundary_conditions = on_architecture(architecture, boundary_conditions)
+    end
+
     close(file)
 
     LX, LY, LZ = Location
diff --git a/src/OutputReaders/set_field_time_series.jl b/src/OutputReaders/set_field_time_series.jl
index 577082782e..c026dc0b5f 100644
--- a/src/OutputReaders/set_field_time_series.jl
+++ b/src/OutputReaders/set_field_time_series.jl
@@ -36,6 +36,7 @@ function set!(fts::InMemoryFTS, path::String=fts.path, name::String=fts.name)
 
         # Note: use the CPU for this step
         field_n = Field(location(fts), path, name, file_iter,
+                        grid = on_architecture(CPU(), fts.grid),    
                         architecture = cpu_architecture(arch),
                         indices = fts.indices,
                         boundary_conditions = fts.boundary_conditions)