JuliaDataCubes / YAXArrays.jl

Yet Another XArray-like Julia package
https://juliadatacubes.github.io/YAXArrays.jl/
Other
101 stars 17 forks source link

Shape of output Cube - regridding #254

Closed Balinus closed 1 year ago

Balinus commented 1 year ago

Hello!

I am trying to understand the following behaviour. I am regridding a Cube onto new coordinates (coords taken from another Cube). When I developed the code outside of the mapCube feature, the resulting dimensions are coherent with the expected sizes. When I transfered the code to use it as mapCube approach, I do not get the same dimensions sizes. I'm sure it's trivial and possibly made an error somewhere, but I can't see it!

Here's a MWE mimicking an interpolation for a single raster.

# template 2 avec lon-lat
using Interpolations

xg1,yg1 = range(-80,stop=-55, length=101),range(40,stop=65, length=101) # source grid
data1 = [sum(xy) for xy=Iterators.product(xg1,yg1)] # toy data for timestep #1 with shape 101x101

# Interpolation struct
itp = interpolate((xg1,yg1), data1, Gridded(Linear()))
etpf = extrapolate(itp, Flat())

# Destination grid
xg2,yg2 = range(-82,stop=-50, length=321),range(42,stop=62, length=201)
coords = Iterators.product(xg2,yg2)

# Doing the interpolations on new grid
data2 = (c->etpf(c...)).(coords) # data2 is a 321x201 matrix

Now, here's the example using Cubes

Functions

function regrid(xout, xin; xg1, yg1, coords)

    xout = Array{typeof(xin[1])}(undef, size(coords))# print(size(xout))

    # print(size(xout))  # gives correct dimensions

    itp = interpolate((xg1,yg1), reverse(xin, dims=1), Gridded(Linear()))
    etpf = extrapolate(itp, Flat())

    # print(size((c->etpf(c...)).(coords))) # gives the correct shape when printing    

    return xout .= (c->etpf(c...)).(coords) # returns "incorrect shape"

end

function regrid(source::YAXArray, dest::YAXArray)

    # harcoded spatial subset to avoid borders "problem" (extrapolation)
    source_spatial = source[longitude = minimum(dest.longitude)..maximum(dest.longitude), latitude=minimum(dest.latitude)..maximum(dest.latitude)]
    dest_spatial = dest[longitude = minimum(source_spatial.longitude)..maximum(source_spatial.longitude), latitude=minimum(source_spatial.latitude)..maximum(source_spatial.latitude)]

    # Grille source
    xg1,yg1 = collect(source_spatial.longitude), reverse(collect(source_spatial.latitude))

    # Grille destination
    xg2,yg2 = collect(dest_spatial.longitude), reverse(collect(dest_spatial.latitude))
    coords = Iterators.product(xg2,yg2)

    # print(size(coords)) # gives correct dimensions

    # Dimensions   
    indims = InDims("longitude", "latitude")
    outdims = OutDims("longitude", "latitude")

    mapCube(regrid, source_spatial; xg1=xg1, yg1=yg1, coords=coords, indims=indims, outdims=outdims)

end 

Cubes


# Create Cubes
axlist = [
    RangeAxis("number", 1:50),
    RangeAxis("longitude", xg1),
    RangeAxis("latitude", reverse(yg1)),    # doing reverse here to mimick my own data and my functions that takes care of reversed data
    ]

data =  rand(50, length(xg1), length(yg1))
source = YAXArray(axlist, data)

axlist = [
    RangeAxis("longitude", xg2),
    RangeAxis("latitude", reverse(yg2)),   # doing reverse here to mimick my own data and my functions that takes care of reversed data
    ]

data = rand(length(xg2), length(yg2))
dest = YAXArray(axlist, data)

Regrid

regrid(source,dest)

YAXArray with the following dimensions
longitude           Axis with 101 Elements from -80.0 to -55.0
latitude            Axis with 81 Elements from 62.0 to 42.0
number              Axis with 50 Elements from 1 to 50
Total size: 3.12 MB
felixcremer commented 1 year ago

When you specify your outdims as longitude and latitude it takes the axes from the input cube. You would have to constrict new lat and Lon axes with the values from coords.

Balinus commented 1 year ago

Oh, I see! Not sure how I should constrict the new lat and lon axes with this mapCube approach? Is there a way to pass the second cube and to link the lat-lon axes from this second cube? (the cube named dest in this case).

Cheers!

felixcremer commented 1 year ago

If you do

outdims = OutDims(axlist)

with the axlist that you use to define the dest cube than this should work.

Balinus commented 1 year ago

Great! I will test later today. Cheers!

Balinus commented 1 year ago

It worked, thanks! Here's the code in case it interest someone. Lots of reverse, but Interpolations only accept ascending index order.

function regrid_cube(xout, xin; xg1, yg1, coords)

    # Interpolation struct
    itp = interpolate((xg1,yg1), reverse(xin, dims=1), Gridded(Linear()))    
    etpf = extrapolate(itp, Flat())

    return xout .= reverse((c->etpf(c...)).(coords), dims=1)

end

"""
    regrid_cube(cube::YAXArray, dest::YAXArray)

Regrid des données de cube (prévisions saisonnières) vers grille dest (ERA5)
"""
function regrid_cube(source::YAXArray, dest::YAXArray)

    # subset "hardcodé" pour éviter effet de bords
    source_spatial = source[longitude = minimum(dest.longitude)..maximum(dest.longitude), latitude=minimum(dest.latitude)..maximum(dest.latitude)]
    dest_spatial = dest[longitude = minimum(source_spatial.longitude)..maximum(source_spatial.longitude), latitude=minimum(source_spatial.latitude)..maximum(source_spatial.latitude)]

    # Grille source
    xg1,yg1 = collect(source_spatial.longitude), reverse(collect(source_spatial.latitude))

    # Grille destination
    xg2,yg2 = collect(dest_spatial.longitude), reverse(collect(dest_spatial.latitude))
    coords = Iterators.product(xg2,yg2)

    # Dimensions
    indims = InDims("longitude","latitude")    
    outdims = OutDims(RangeAxis("longitude", xg2), RangeAxis("latitude", reverse(yg2)))

    mapCube(regrid_cube, source_spatial; xg1=xg1, yg1=yg1, coords=coords, indims=indims, outdims=outdims)

end