Add Var.are_compatible for FlatVar and Metadata

NEWS.md

@@ -45,6 +45,7 @@ Makie.save("surface_plot.png", fig)
 - The number of dimensions of a `OutputVar` can be accessed with `ndims`.
 - You can drop dimensions of size 1 with `dropdims`.
 - `==` and `isequal` are now defined for `OutputVar`, `FlatVar`, and `Metadata`.
+- Add `arecompatible` for `FlatVar` and `Metadata`.
 
 ## Bug fixes
 

docs/src/api.md

@@ -58,7 +58,7 @@ Var.variance_time
 Var.window
 Var.select
 Var.view_select
-Var.arecompatible
+Var.arecompatible(x::OutputVar, y::OutputVar)
 Var.shift_longitude
 Var.center_longitude!
 Var.times
@@ -165,6 +165,8 @@ Var.flatten
 Var.unflatten
 Var.flatten_dim_order
 Var.flattened_length
+Var.arecompatible(x::ClimaAnalysis.Var.FlatVar, y::ClimaAnalysis.Var.FlatVar)
+Var.arecompatible(x::ClimaAnalysis.Var.Metadata, y::ClimaAnalysis.Var.Metadata)
 ```
 
 ## Leaderboard

src/flat.jl

@@ -150,17 +150,6 @@ function flatten(var::OutputVar, metadata::Metadata)
         "Dimensions in var ($(dim_names(var))) is not the same as the dimensions in the metadata ($(keys(metadata.dims)))",
     )
 
-    # Instead of trying to determine whether dates or times should be used, we try dates
-    # first and if it does not work, then use the time dimension instead
-    function dates_or_times(var)
-        temporal_dim = try
-            dates(var)
-        catch
-            times(var)
-        end
-        return temporal_dim
-    end
-
     # Check the dimension arrays are the same
     for var_dim_name in dim_names(var)
         md_dim_name = find_corresponding_dim_name_in_var(var_dim_name, metadata)
@@ -170,7 +159,7 @@ function flatten(var::OutputVar, metadata::Metadata)
                     "Dimensions in var ($var_dim_name) and metadata ($md_dim_name) are not the same",
                 )
         else
-            dates_or_times(var) == dates_or_times(metadata) || error(
+            _dates_or_times(var) == _dates_or_times(metadata) || error(
                 "The dates or times between var and metadata are not the same",
             )
         end
@@ -271,10 +260,95 @@ function unflatten(metadata::Metadata, data::AbstractVector)
     )
 end
 
+"""
+    arecompatible(x::FlatVar, y::FlatVar; ignore_dims = ())
+
+Return whether two `FlatVar` are defined on the same physical space, with values
+dropped at the same coordinates.
+
+See the documentation for
+[`arecompatible(x::ClimaAnalysis.Var.Metadata, y::ClimaAnalysis.Var.Metadata)`](@ref) for
+more information.
+"""
+function arecompatible(x::FlatVar, y::FlatVar; ignore_dims = ())
+    return arecompatible(x.metadata, y.metadata; ignore_dims = ignore_dims)
+end
+
+"""
+    arecompatible(x::Metadata, y::Metadata; ignore_dims = ())
+
+Return whether two `Metadata` are compatible.
+
+Two `Metadata` are compatible when the `FlatVar`s built from them are defined on the same
+physical space, with values dropped at the same coordinates.
+
+This is accomplished by comparing `dims` and `dim_attributes` (the latter because they might
+contain information about the units), the coordinates at which the values are dropped, and
+the order that the dimensions are flattened.
+
+The keyword argument `ignore_dims` is a dimension name or collection of dimension names
+whose coordinate values and dimension units are not compared between the `Metadata`s.
+"""
+function arecompatible(x::Metadata, y::Metadata; ignore_dims = ())
+    ignore_dims = ignore_dims isa AbstractString ? (ignore_dims,) : ignore_dims
+    ignored_names = Set(conventional_dim_name.(ignore_dims))
+
+    # Check the order of the dimensions when flattening are the same and the
+    # type of dimensions
+    conventional_dim_name.(flatten_dim_order(x)) ==
+    conventional_dim_name.(flatten_dim_order(y)) || return false
+
+    for x_dim_name in dim_names(x)
+        conventional_dim_name(x_dim_name) in ignored_names && continue
+
+        # Check coordinate values of the corresponding dimensions are the same
+        y_dim_name = find_corresponding_dim_name_in_var(x_dim_name, y)
+        if conventional_dim_name(x_dim_name) != "time"
+            length(x.dims[x_dim_name]) == length(y.dims[y_dim_name]) ||
+                return false
+            isapprox(x.dims[x_dim_name], y.dims[y_dim_name]) || return false
+        else
+            _dates_or_times(x) == _dates_or_times(y) || return false
+        end
+
+        # Check the units of the corresponding dimensions are the same
+        x_dim_units = dim_units(x, x_dim_name)
+        y_dim_units = dim_units(y, y_dim_name)
+        x_dim_units == y_dim_units || return false
+
+        # Warn if dimension units are empty
+        isempty(x_dim_units) &&
+            @warn "Missing units for dimension $x_dim_name in FlatVar/Metadata with short name $(short_name(x))"
+        isempty(y_dim_units) &&
+            @warn "Missing units for dimension $y_dim_name in FlatVar/Metadata with short name $(short_name(y))"
+    end
+
+    # Check the drop_mask are the same, since we care about which coordinates
+    # are kept and dropped
+    return x.drop_mask == y.drop_mask
+end
+
+"""
+    _dates_or_times(var_or_metadata)
+
+Helper function to return the coordinate values of the time dimension of `var_or_metadata`,
+preferring dates over times.
+
+This function does not check if `var_or_metadata` has a time dimension.
+"""
+function _dates_or_times(var_or_metadata)
+    temporal_dim = try
+        dates(var_or_metadata)
+    catch
+        times(var_or_metadata)
+    end
+    return temporal_dim
+end
+
 """
     flatten_dim_order(var::FlatVar)
 
-Return the order of the dimensions before flattening `data`.
+Return the order of the dimensions when flattening the `data` of `var`.
 """
 function flatten_dim_order(var::FlatVar)
     return flatten_dim_order(var.metadata)
@@ -283,7 +357,8 @@ end
 """
     flatten_dim_order(metadata::Metadata)
 
-Return the order of the dimensions before flattening `data`.
+Return the order of the dimensions when flattening the `data` of `var` as determined from
+the `metadata`.
 """
 function flatten_dim_order(metadata::Metadata)
     return metadata.ordered_dims

test/test_flat.jl

@@ -408,6 +408,140 @@ end
     test_flatten(nan_var, nan_var, perms_3d)
 end
 
+@testset "Are compatible" begin
+    lat = [-90.0, -30.0, 30.0, 90.0]
+    lon = [-60.0, -30.0, 0.0, 30.0, 60.0]
+    time = [0.0, 1.0, 5.0]
+    var =
+        TemplateVar() |>
+        add_dim("lat", lat, units = "degrees") |>
+        add_dim("lon", lon, units = "degrees") |>
+        add_dim("time", time, units = "seconds") |>
+        add_attribs(long_name = "hi", start_date = "2010-12-1") |>
+        initialize
+
+    var2 = var + 10
+
+    function are_flatten_compatible(
+        x::ClimaAnalysis.OutputVar,
+        y::ClimaAnalysis.OutputVar;
+        ignore_dims = (),
+        perm = ("longitude", "latitude", "time", "pfull"),
+        other_perm = nothing,
+        ignore_nan = true,
+    )
+        isnothing(other_perm) && (other_perm = perm)
+        x = ClimaAnalysis.flatten(x; ignore_nan, dims = perm)
+        y = ClimaAnalysis.flatten(y; ignore_nan, dims = other_perm)
+        flat_var_compatible = ClimaAnalysis.arecompatible(x, y; ignore_dims)
+        metadata_compatible =
+            ClimaAnalysis.arecompatible(x.metadata, y.metadata; ignore_dims)
+        return flat_var_compatible && metadata_compatible
+    end
+
+    # Test with modified data
+    @test are_flatten_compatible(var, var2)
+
+    # Test with permuted dimensions
+    perms_3d = (
+        ("lon", "lat", "time"),
+        ("lat", "lon", "time"),
+        ("lat", "time", "lon"),
+        ("lon", "time", "lat"),
+        ("time", "lon", "lat"),
+        ("time", "lat", "lon"),
+    )
+    for perm in perms_3d
+        @test are_flatten_compatible(var, var2; perm)
+    end
+
+    # Not compatible with dates and floating point times
+    no_dates_var = deepcopy(var)
+    delete!(no_dates_var.attributes, "start_date")
+    @test !are_flatten_compatible(var, no_dates_var)
+
+    # Test with ignore_dims keyword argument
+    @test are_flatten_compatible(var, no_dates_var, ignore_dims = "time")
+    @test are_flatten_compatible(
+        var,
+        no_dates_var,
+        ignore_dims = ["t", "latitude"],
+    )
+
+    # Different dimensions but same number
+    pfull = [0.0, 1.0, 5.0]
+    no_time_var =
+        TemplateVar() |>
+        add_dim("lat", lat, units = "degrees") |>
+        add_dim("lon", lon, units = "degrees") |>
+        add_dim("pfull", time, units = "seconds") |>
+        add_attribs(long_name = "hi") |>
+        initialize
+    @test !are_flatten_compatible(var, no_time_var)
+
+    # Extra dimension
+    more_dim_var =
+        TemplateVar() |>
+        add_dim("lat", lat, units = "degrees") |>
+        add_dim("lon", lon, units = "degrees") |>
+        add_dim("time", time, units = "seconds") |>
+        add_dim("pfull", time, units = "seconds") |>
+        add_attribs(long_name = "hi") |>
+        initialize
+    @test !are_flatten_compatible(var, more_dim_var)
+
+    # Fewer dimension
+    less_dim_var =
+        TemplateVar() |>
+        add_dim("lat", lat, units = "degrees") |>
+        add_dim("lon", lon, units = "degrees") |>
+        add_attribs(long_name = "hi") |>
+        initialize
+    @test !are_flatten_compatible(var, less_dim_var)
+
+    # Drop mask is different
+    var = ClimaAnalysis.remake(var, data = collect(Float64.(var.data)))
+    nan_var1 = deepcopy(var)
+    nan_var2 = deepcopy(var)
+    nan_var1.data[begin] = NaN
+    nan_var2.data[end] = NaN
+    @test !are_flatten_compatible(nan_var1, nan_var2)
+    @test are_flatten_compatible(nan_var1, nan_var2, ignore_nan = false)
+
+    # Same dimensions but different length
+    windowed_var = ClimaAnalysis.window(var, "time", left = 0.0, right = 1.0)
+    @test !are_flatten_compatible(var, windowed_var)
+
+    # Same dimensions but different dimension values
+    shifted_lon_var = ClimaAnalysis.shift_longitude(var, 0.0, 360.0)
+    @test !are_flatten_compatible(var, shifted_lon_var)
+
+    # Different dimension units
+    rad_var = deepcopy(var)
+    ClimaAnalysis.set_dim_units!(rad_var, "lon", "radians")
+    @test !are_flatten_compatible(var, rad_var)
+
+    # Different flatten order
+    @test !are_flatten_compatible(
+        var,
+        var;
+        ignore_dims = (),
+        perm = ("longitude", "latitude", "time"),
+        other_perm = ("time", "longitude", "latitude"),
+        ignore_nan = true,
+    )
+
+    # Warn on empty units
+    missing_dim_units_var = deepcopy(var)
+    ClimaAnalysis.set_dim_units!(missing_dim_units_var, "longitude", "")
+    @test_logs (:warn, r"Missing units for dimension lon") match_mode = :any ClimaAnalysis.arecompatible(
+        ClimaAnalysis.flatten.([
+            missing_dim_units_var,
+            missing_dim_units_var,
+        ])...,
+    )
+end
+
 @testset "Extracting dimensions, units, and names for FlatVar" begin
     lat = collect(range(-89.5, 89.5, 3))
     lon = collect(range(-179.5, 179.5, 4))