Add different download directories for JRA55 and Bathymetry

src/DataWrangling/ECCO.jl

@@ -139,7 +139,7 @@ function ecco_field(metadata::ECCOMetadata;
 
     download_dataset!(metadata)
 
-    ds = Dataset(filename)
+    ds = Dataset(joinpath(metadata.dir, filename))
     if variable_is_three_dimensional(metadata)
         data = ds[shortname][:, :, :, 1]
         # The surface layer in three-dimensional ECCO fields is at `k = 1`

src/DataWrangling/ecco_metadata.jl

@@ -1,9 +1,15 @@
 using CFTime
 using Dates
+using Scratch
 
 import Oceananigans.Fields: set!
 import Base
 
+download_ECCO_cache::String = ""
+function __init__()
+    global download_ECCO_cache = @get_scratch!("ECCO")
+end
+
 struct ECCO2Monthly end
 struct ECCO2Daily end
 struct ECCO4Monthly end
@@ -16,6 +22,7 @@ struct ECCOMetadata{D, V}
     name  :: Symbol
     dates :: D
     version :: V
+    dir :: String
 end
 
 Base.show(io::IO, metadata::ECCOMetadata) = 
@@ -27,18 +34,21 @@ Base.show(io::IO, metadata::ECCOMetadata) =
 # The default is the ECCO2Daily dataset at 1993-01-01.
 function ECCOMetadata(name::Symbol; 
                       date = DateTimeProlepticGregorian(1993, 1, 1),
-                   version = ECCO2Daily()) 
+                   version = ECCO2Daily(),
+                       dir = download_ECCO_cache) 
 
-    return ECCOMetadata(name, date, version) 
+    return ECCOMetadata(name, date, version, dir) 
 end
 
+ECCOMetadata(name::Symbol, date, version=ECCO4Monthly(); dir = download_ECCO_cache) = ECCOMetadata(name, date, version, dir)
+
 # Treat ECCOMetadata as an array to allow iteration over the dates.
-Base.getindex(metadata::ECCOMetadata, i::Int) = @inbounds ECCOMetadata(metadata.name, metadata.dates[i], metadata.version)
+Base.getindex(metadata::ECCOMetadata, i::Int) = @inbounds ECCOMetadata(metadata.name, metadata.dates[i], metadata.version, metadata.dir)
 Base.length(metadata::ECCOMetadata)       = length(metadata.dates)
 Base.eltype(metadata::ECCOMetadata)       = Base.eltype(metadata.dates)
-Base.first(metadata::ECCOMetadata)        = @inbounds ECCOMetadata(metadata.name, metadata.dates[1], metadata.version)
-Base.last(metadata::ECCOMetadata)         = @inbounds ECCOMetadata(metadata.name, metadata.dates[end], metadata.version)
-Base.iterate(metadata::ECCOMetadata, i=1) = (@inline; (i % UInt) - 1 < length(metadata) ? (@inbounds ECCOMetadata(metadata.name, metadata.dates[i], metadata.version), i + 1) : nothing)
+Base.first(metadata::ECCOMetadata)        = @inbounds ECCOMetadata(metadata.name, metadata.dates[1], metadata.version, metadata.dir)
+Base.last(metadata::ECCOMetadata)         = @inbounds ECCOMetadata(metadata.name, metadata.dates[end], metadata.version, metadata.dir)
+Base.iterate(metadata::ECCOMetadata, i=1) = (@inline; (i % UInt) - 1 < length(metadata) ? (@inbounds ECCOMetadata(metadata.name, metadata.dates[i], metadata.version, metadata.dir), i + 1) : nothing)
 
 Base.axes(metadata::ECCOMetadata{<:AbstractCFDateTime})    = 1
 Base.first(metadata::ECCOMetadata{<:AbstractCFDateTime})   = metadata
@@ -146,6 +156,7 @@ function download_dataset!(metadata::ECCOMetadata;
 
     username = get(ENV, "ECCO_USERNAME", nothing)
     password = get(ENV, "ECCO_PASSWORD", nothing)
+    dir = metadata.dir
 
     for data in metadata
         filename  = metadata_filename(data)
@@ -164,7 +175,7 @@ function download_dataset!(metadata::ECCOMetadata;
                 fileurl = joinpath(url, shortname, year, filename)
             end
 
-            cmd = `wget --http-user=$(username) --http-passwd=$(password) $(fileurl)`
+            cmd = `wget --http-user=$(username) --http-passwd=$(password)  --directory-prefix=$(dir) $(fileurl)`
 
             run(cmd)
         end

src/DataWrangling/ecco_restoring.jl

@@ -52,7 +52,7 @@ function set!(fts::ECCONetCDFFTS, path::ECCOMetadata=fts.path, name::String=fts.
         metadata = @inbounds path[t] 
 
         arch = architecture(fts)
-        f = inpainted_ecco_field(metadata; architecture = arch)
+        f = inpainted_ecco_field(metadata; architecture=arch)
         if on_native_grid(backend)
             set!(fts[t], f)
         else

test/runtests_setup.jl

@@ -15,3 +15,6 @@ using ClimaOcean
 
 gpu_test = parse(Bool, get(ENV, "GPU_TEST", "false"))
 test_architectures = gpu_test ? [GPU()] : [CPU()]
+JRA55_data_directory = gpu_test ? "GPU_JRA55_data" : "CPU_JRA55_data"
+bathymetry_data_directory = gpu_test ? "GPU_Bathymetry_data" : "CPU_Bathymetry_data"
+ECCO_data_directory = gpu_test ? "GPU_ECCO_data" : "CPU_ECCO_data"  

test/test_bathymetry.jl

@@ -6,14 +6,21 @@ using Statistics
 @testset "Availability of Bathymetry" begin
     @info "Testing Bathymetry utils..."
     for arch in test_architectures
+
+        # Make data directory if it doesn't exist
+        try
+            mkdir(bathymetry_data_directory)
+        catch
+        end
+
         grid = LatitudeLongitudeGrid(arch;
                                      size = (100, 100, 10), 
                                      longitude = (0, 100), 
                                      latitude = (0, 50),
                                      z = (-6000, 0))
 
         # Test that remove_minor_basins!(Z, Inf) does nothing
-        control_bottom_height = regrid_bathymetry(grid)        
+        control_bottom_height = regrid_bathymetry(grid; dir=bathymetry_data_directory)        
         bottom_height = deepcopy(control_bottom_height)
         @test_throws ArgumentError remove_minor_basins!(bottom_height, Inf)
 
@@ -52,10 +59,13 @@ using Statistics
                                      latitude = (-10, 50),
                                      z = (-6000, 0))
 
-        control_bottom_height = regrid_bathymetry(grid)
-        interpolated_bottom_height = regrid_bathymetry(grid; interpolation_passes = 100)
+        control_bottom_height = regrid_bathymetry(grid; dir=bathymetry_data_directory)
+        interpolated_bottom_height = regrid_bathymetry(grid; interpolation_passes=100, dir=bathymetry_data_directory)
 
         # Testing that multiple passes do not change the solution when refining the grid
         @test parent(control_bottom_height) == parent(interpolated_bottom_height)
+
+        # Remove the data directory
+        rm(bathymetry_data_directory; recursive=true)
     end
 end 

test/test_ecco.jl

@@ -8,6 +8,12 @@ using Oceananigans.Grids: topology
 using CFTime
 using Dates
 
+# Make data directory if it doesn't exist
+try
+    mkdir(ECCO_data_directory)
+catch
+end
+
 @testset "ECCO fields utilities" begin
     for arch in test_architectures
         A = typeof(arch)
@@ -17,14 +23,14 @@ using Dates
         end_date = DateTimeProlepticGregorian(1993, 4, 1)
         dates = start_date : Month(1) : end_date
 
-        temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly())
+        temperature = ECCOMetadata(:temperature, dates; dir=ECCO_data_directory)
         t_restoring = ECCO_restoring_forcing(temperature; timescale = 1000.0)
 
         ecco_fts = t_restoring.func.ecco_fts
 
         for metadata in temperature
             temperature_filename = metadata_filename(metadata)
-            @test isfile(temperature_filename)
+            @test isfile(joinpath(ECCO_data_directory, temperature_filename))
         end
 
         @test ecco_fts isa FieldTimeSeries
@@ -45,7 +51,9 @@ end
     for arch in test_architectures
         grid  = LatitudeLongitudeGrid(size = (10, 10, 10), latitude = (-60, -40), longitude = (10, 15), z = (-200, 0))
         field = CenterField(grid)
-        set!(field, ECCOMetadata(:temperature)) 
-        set!(field, ECCOMetadata(:salinity))
+        set!(field, ECCOMetadata(:temperature; dir=ECCO_data_directory)) 
+        set!(field, ECCOMetadata(:salinity; dir=ECCO_data_directory))
     end 
 end
+
+rm(ECCO_data_directory; recursive=true)

test/test_jra55.jl

@@ -11,10 +11,16 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere
         test_name = :downwelling_shortwave_radiation
         time_indices = 1:3
 
+        # Make data directory if it doesn't exist
+        try
+            mkdir(JRA55_data_directory)
+        catch
+        end
+
         # This should download a file called "RYF.rsds.1990_1991.nc"
-        jra55_fts = JRA55_field_time_series(test_name; architecture=arch, time_indices)
+        jra55_fts = JRA55_field_time_series(test_name; architecture=arch, time_indices, dir=JRA55_data_directory)
 
-        test_filename = joinpath(download_jra55_cache, "RYF.rsds.1990_1991.nc")
+        test_filename = joinpath(JRA55_data_directory, "RYF.rsds.1990_1991.nc")
 
         @test jra55_fts isa FieldTimeSeries
         @test jra55_fts.grid isa LatitudeLongitudeGrid
@@ -39,7 +45,8 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere
         in_memory_jra55_fts = JRA55_field_time_series(test_name;
                                                       time_indices,
                                                       architecture = arch,
-                                                      backend = InMemory(2))
+                                                      backend = InMemory(2),
+                                                      dir = JRA55_data_directory)
 
         @test in_memory_jra55_fts isa FieldTimeSeries
 
@@ -101,13 +108,16 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere
         #####
 
         backend    = JRA55NetCDFBackend(2) 
-        atmosphere = JRA55_prescribed_atmosphere(arch; backend, include_rivers_and_icebergs=false)
+        atmosphere = JRA55_prescribed_atmosphere(arch; backend, include_rivers_and_icebergs=false, dir=JRA55_data_directory)
         @test atmosphere isa PrescribedAtmosphere
         @test isnothing(atmosphere.auxiliary_freshwater_flux)
 
-        atmosphere = JRA55_prescribed_atmosphere(arch; backend, include_rivers_and_icebergs=true)
+        atmosphere = JRA55_prescribed_atmosphere(arch; backend, include_rivers_and_icebergs=true, dir=JRA55_data_directory)
         @test haskey(atmosphere.auxiliary_freshwater_flux, :rivers)
         @test haskey(atmosphere.auxiliary_freshwater_flux, :icebergs)
+
+        # Remove the data directory
+        rm(JRA55_data_directory; recursive=true)
     end 
 end