def test_padding_fill(self, boundary_width, fill_value): ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords) data = ds.tracer # iterate over all axes expected = data.copy(deep=True) for ax, widths in boundary_width.items(): dim = grid._get_dims_from_axis( data, ax)[0] # ? not entirely sure why this is a list expected = expected.pad({dim: widths}, "constant", constant_values=fill_value) # we intentionally strip all coords from padded arrays expected = _strip_all_coords(expected) result = pad( data, grid, boundary="fill", boundary_width=boundary_width, fill_value=fill_value, other_component=None, ) xr.testing.assert_allclose(expected, result)
def test_padding_mixed(self, boundary_width): ds, coords, _ = datasets_grid_metric("C") grid = Grid(ds, coords=coords) data = ds.tracer axis_padding_mapping = {"X": "periodic", "Y": "extend"} method_mapping = { "periodic": "wrap", "extend": "edge", } # iterate over all axes expected = data.copy(deep=True) for ax, widths in boundary_width.items(): dim = grid._get_dims_from_axis( data, ax)[0] # ? not entirely sure why this is a list expected = expected.pad({dim: widths}, method_mapping[axis_padding_mapping[ax]]) # we intentionally strip all coords from padded arrays expected = _strip_all_coords(expected) result = pad( data, grid, boundary=axis_padding_mapping, boundary_width=boundary_width, fill_value=None, other_component=None, ) xr.testing.assert_allclose(expected, result)
def test_face_connections_right_left_same_axis(self, boundary_width, ds_faces, fill_value): face_connections = { "face": { 0: { "X": (None, (1, "X", False)) }, 1: { "X": ((0, "X", False), None) }, } } grid = Grid(ds_faces, face_connections=face_connections) data = ds_faces.data_c # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) # # restrict data here, so its easier to see the output # data = data.isel(y=slice(0, 2), x=slice(0, 2)) data = data.reset_coords(drop=True).reset_index(data.dims, drop=True) # prepad the arrays face_0_padded, face_1_padded = _prepad_right_left_same_axis( data, boundary_width, fill_value) # then simply add the corresponding slice to each face according to the connection face_0_expected = xr.concat( [ face_0_padded, face_1_padded.isel(x=slice(0, boundary_width["X"][1])) ], dim="x", ) face_1_expected = xr.concat( [ face_0_padded.isel(x=slice( -boundary_width["X"][0], None if boundary_width["X"][0] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )), face_1_padded, ], dim="x", ) expected = xr.concat([face_0_expected, face_1_expected], dim="face") result = pad( data, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component=None, ) xr.testing.assert_allclose(result, expected)
def test_vector_face_connections_coord_padding(self, boundary_width, ds_faces, fill_value, boundary): # make sure that the complex padding acts like xarray.pad when it comes to dimension coordinates face_connections = { "face": { 0: { "X": (None, (1, "Y", True)) }, 1: { "Y": (None, (0, "X", True)) }, } } grid = Grid(ds_faces, face_connections=face_connections) u = ds_faces.u v = ds_faces.v # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) padded_complex = pad( {"X": u}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"Y": v}, ) padded_simple = u.pad( { "xl": boundary_width["X"], "y": boundary_width["Y"] }, "constant", constant_values=fill_value, ) # TODO: add constant coord padding here once implemented in xarray. for di in u.dims: assert (di in padded_simple.coords and di in padded_complex.coords ) or (di not in padded_simple.coords and di not in padded_complex.coords) if di in padded_simple.coords: xr.testing.assert_allclose(padded_complex[di], padded_simple[di])
def test_vector_face_connections_right_right_swap_axis( self, boundary_width, ds_faces, fill_value): face_connections = { "face": { 0: { "X": (None, (1, "Y", True)) }, 1: { "Y": (None, (0, "X", True)) }, } } grid = Grid(ds_faces, face_connections=face_connections) u = ds_faces.u v = ds_faces.v # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) # # restrict data here, so its easier to see the output u = u.reset_coords(drop=True).reset_index(u.dims, drop=True) v = v.reset_coords(drop=True).reset_index(v.dims, drop=True) ( u_face_0_padded, u_face_1_padded, u_face_0_padded_swapped, u_face_1_padded_swapped, ) = _prepad_right_right_swap_axis(u, boundary_width, fill_value, x="xl", y="y") ( v_face_0_padded, v_face_1_padded, v_face_0_padded_swapped, v_face_1_padded_swapped, ) = _prepad_right_right_swap_axis(v, boundary_width, fill_value, x="x", y="yl") # Put together the additions for each face u_face_0_addition = v_face_1_padded_swapped.isel(yl=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )).rename({ "x": "xl", "yl": "y" }) u_face_0_addition = -u_face_0_addition.isel(y=slice(None, None, -1)) u_face_0_addition = _maybe_swap_dimension_names( u_face_0_addition, "y", "xl") u_face_1_addition = v_face_0_padded_swapped.isel(x=slice( -boundary_width["Y"][1], None if boundary_width["Y"][1] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )).rename({ "x": "xl", "yl": "y" }) u_face_1_addition = u_face_1_addition.isel(xl=slice(None, None, -1)) u_face_1_addition = _maybe_swap_dimension_names( u_face_1_addition, "y", "xl") # now v (this one needs a sign change) v_face_0_addition = u_face_1_padded_swapped.isel(y=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )).rename({ "xl": "x", "y": "yl" }) # Tangential flip (v DOES need sign change) v_face_0_addition = v_face_0_addition.isel(yl=slice(None, None, -1)) v_face_0_addition = _maybe_swap_dimension_names( v_face_0_addition, "yl", "x") v_face_1_addition = u_face_0_padded_swapped.isel(xl=slice( -boundary_width["Y"][1], None if boundary_width["Y"][1] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )).rename({ "xl": "x", "y": "yl" }) v_face_1_addition = -v_face_1_addition.isel(x=slice(None, None, -1)) v_face_1_addition = _maybe_swap_dimension_names( v_face_1_addition, "yl", "x") # then simply add the corresponding slice to each face according to the connection u_face_0_expected = xr.concat( [u_face_0_padded, u_face_0_addition], dim="xl", ) u_face_1_expected = xr.concat( [u_face_1_padded, u_face_1_addition], dim="y", ) v_face_0_expected = xr.concat( [v_face_0_padded, v_face_0_addition], dim="x", ) v_face_1_expected = xr.concat( [v_face_1_padded, v_face_1_addition], dim="yl", ) u_expected = xr.concat([u_face_0_expected, u_face_1_expected], dim="face") v_expected = xr.concat([v_face_0_expected, v_face_1_expected], dim="face") u_result = pad( {"X": u}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"Y": v}, ) v_result = pad( {"Y": v}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"X": u}, ) xr.testing.assert_allclose(u_result, u_expected) xr.testing.assert_allclose(v_result, v_expected)
def test_vector_face_connections_right_right_same_axis( self, boundary_width, ds_faces, fill_value): face_connections = { "face": { 0: { "X": (None, (1, "X", True)) }, 1: { "X": (None, (0, "X", True)) }, } } grid = Grid(ds_faces, face_connections=face_connections) u = ds_faces.u v = ds_faces.v # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) # # restrict data here, so its easier to see the output # data = data.isel(y=slice(0, 2), x=slice(0, 2)) u = u.reset_coords(drop=True).reset_index(u.dims, drop=True) v = v.reset_coords(drop=True).reset_index(v.dims, drop=True) u_face_0_padded, u_face_1_padded = _prepad_right_right_same_axis( u, boundary_width, fill_value, x="xl", y="y") v_face_0_padded, v_face_1_padded = _prepad_right_right_same_axis( v, boundary_width, fill_value, x="x", y="yl", ) # Slice the appropriate portion of the source face to concat u_face_0_addition = u_face_1_padded.isel(xl=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, )) u_face_1_addition = u_face_0_padded.isel(xl=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, )) v_face_0_addition = v_face_1_padded.isel(x=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, )) v_face_1_addition = v_face_0_padded.isel(x=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, )) # do a parallel flip since this connection is reversed (and sign change for u) u_face_0_addition = -u_face_0_addition.isel(xl=slice(None, None, -1)) u_face_1_addition = -u_face_1_addition.isel(xl=slice(None, None, -1)) v_face_0_addition = v_face_0_addition.isel(x=slice(None, None, -1)) v_face_1_addition = v_face_1_addition.isel(x=slice(None, None, -1)) # then simply add the corresponding slice to each face according to the connection u_face_0_expected = xr.concat( [u_face_0_padded, u_face_0_addition], dim="xl", ) u_face_1_expected = xr.concat( [u_face_1_padded, u_face_1_addition], dim="xl", ) v_face_0_expected = xr.concat( [v_face_0_padded, v_face_0_addition], dim="x", ) v_face_1_expected = xr.concat( [v_face_1_padded, v_face_1_addition], dim="x", ) u_expected = xr.concat([u_face_0_expected, u_face_1_expected], dim="face") v_expected = xr.concat([v_face_0_expected, v_face_1_expected], dim="face") # test u u_result = pad( {"X": u}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"Y": v}, ) xr.testing.assert_allclose(u_result, u_expected) # test v v_result = pad( {"Y": v}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"X": u}, ) xr.testing.assert_allclose(v_result, v_expected)
def test_vector_face_connections_right_left_same_axis( self, boundary_width, ds_faces, fill_value): face_connections = { "face": { 0: { "X": (None, (1, "X", False)) }, 1: { "X": ((0, "X", False), None) }, } } grid = Grid(ds_faces, face_connections=face_connections) u = ds_faces.u v = ds_faces.v # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) # # restrict data here, so its easier to see the output # data = data.isel(y=slice(0, 2), x=slice(0, 2)) u = u.reset_coords(drop=True).reset_index(u.dims, drop=True) v = v.reset_coords(drop=True).reset_index(v.dims, drop=True) # prepad the arrays u_face_0_padded, u_face_1_padded = _prepad_right_left_same_axis( u, boundary_width, fill_value, x="xl", y="y") v_face_0_padded, v_face_1_padded = _prepad_right_left_same_axis( v, boundary_width, fill_value, x="x", y="yl", ) # Slice the appropriate portion of the source face to concat u_face_0_addition = u_face_1_padded.isel( xl=slice(0, boundary_width["X"][1])) u_face_1_addition = u_face_0_padded.isel(xl=slice( -boundary_width["X"][0], None if boundary_width["X"][0] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )) v_face_0_addition = v_face_1_padded.isel( x=slice(0, boundary_width["X"][1])) v_face_1_addition = v_face_0_padded.isel(x=slice( -boundary_width["X"][0], None if boundary_width["X"][0] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )) # then simply add the corresponding slice to each face according to the connection u_face_0_expected = xr.concat( [u_face_0_padded, u_face_0_addition], dim="xl", ) u_face_1_expected = xr.concat( [u_face_1_addition, u_face_1_padded], dim="xl", ) v_face_0_expected = xr.concat( [v_face_0_padded, v_face_0_addition], dim="x", ) v_face_1_expected = xr.concat( [v_face_1_addition, v_face_1_padded], dim="x", ) u_expected = xr.concat([u_face_0_expected, u_face_1_expected], dim="face") v_expected = xr.concat([v_face_0_expected, v_face_1_expected], dim="face") # test u u_result = pad( {"X": u}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"Y": v}, ) xr.testing.assert_allclose(u_result, u_expected) # test v v_result = pad( {"Y": v}, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component={"X": u}, ) xr.testing.assert_allclose(v_result, v_expected)
def test_face_connections_right_right_swap_axis(self, boundary_width, ds_faces, fill_value): # set a default for boundary widths boundary_width = {k: boundary_width.get(k, (0, 0)) for k in ["X", "Y"]} face_connections = { "face": { 0: { "X": (None, (1, "Y", True)) }, 1: { "Y": (None, (0, "X", True)) }, } } grid = Grid(ds_faces, face_connections=face_connections) data = ds_faces.data_c # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) # restrict data here, so its easier to see the output data = data.isel(y=slice(0, 3), x=slice(0, 3)) data = data.reset_coords(drop=True).reset_index(data.dims, drop=True) ( face_0_padded, face_1_padded, face_0_padded_swapped, face_1_padded_swapped, ) = _prepad_right_right_swap_axis(data, boundary_width, fill_value) # then simply add the corresponding slice to each face according to the connection # in this case we also need to rename them face_0_addition = face_1_padded_swapped.isel(y=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )) # Same steps for the other face face_1_addition = face_0_padded_swapped.isel(x=slice( -boundary_width["Y"][1], None if boundary_width["Y"][1] > 0 else 0, # this is a bit annoying. if the boundary width on this side is # 0 I want nothing to be padded. but slice(0,None) pads the whole array... )) # Flip both of these along the parallel axis face_0_addition = face_0_addition.isel(y=slice(None, None, -1)) face_1_addition = face_1_addition.isel(x=slice(None, None, -1)) # In this case we need to rename the 'addition' dimensions face_0_addition = _maybe_swap_dimension_names(face_0_addition, "y", "x") face_1_addition = _maybe_swap_dimension_names(face_1_addition, "x", "y") face_0_expected = xr.concat( [face_0_padded, face_0_addition], dim="x", ) face_1_expected = xr.concat( [face_1_padded, face_1_addition], dim="y", ) expected = xr.concat([face_0_expected, face_1_expected], dim="face") result = pad( data, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component=None, ) xr.testing.assert_allclose(result, expected)
def test_face_connections_right_right_same_axis(self, boundary_width, ds_faces, fill_value): face_connections = { "face": { 0: { "X": (None, (1, "X", True)) }, 1: { "X": (None, (0, "X", True)) }, } } grid = Grid(ds_faces, face_connections=face_connections) data = ds_faces.data_c # fill in zeros for y boundary width if not given boundary_width["Y"] = boundary_width.get("Y", (0, 0)) # # restrict data here, so its easier to see the output # data = data.isel(y=slice(0, 2), x=slice(0, 2)) data = data.reset_coords(drop=True).reset_index(data.dims, drop=True) # prepad the arrays face_0_padded, face_1_padded = _prepad_right_right_same_axis( data, boundary_width, fill_value) # Process the padded data face_0_addition = face_1_padded.isel(x=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, )) face_1_addition = face_0_padded.isel(x=slice( -boundary_width["X"][1], None if boundary_width["X"][1] > 0 else 0, )) # do a parallel flip since this connection is reversed face_0_addition = face_0_addition.isel(x=slice(None, None, -1)) face_1_addition = face_1_addition.isel(x=slice(None, None, -1)) # then simply add the corresponding slice to each face according to the connection face_0_expected = xr.concat( [face_0_padded, face_0_addition], dim="x", ) face_1_expected = xr.concat( [face_1_padded, face_1_addition], dim="x", ) expected = xr.concat([face_0_expected, face_1_expected], dim="face") result = pad( data, grid, boundary_width=boundary_width, boundary="fill", fill_value=fill_value, other_component=None, ) xr.testing.assert_allclose(result, expected)