sectionate 0.3.0__py3-none-any.whl

sectionate/tests/test_convergent_transport.py

@@ -0,0 +1,120 @@
+import numpy as np
+import xarray as xr
+import xgcm
+
+# define simple xgcm grid
+def initialize_minimal_outer_grid():
+    xh = np.array([0.5])
+    yh = np.array([0.5])
+    xq = np.array([0., 1.])
+    yq = np.array([0., 1.])
+
+    lon, lat = np.meshgrid(xh, yh)
+    lon_c, lat_c = np.meshgrid(xq, yq)
+    ds = xr.Dataset({}, coords={
+        "xh":xr.DataArray(xh, dims=("xh",)),
+        "yh":xr.DataArray(yh, dims=("yh",)),
+        "xq":xr.DataArray(xq, dims=("xq",)),
+        "yq":xr.DataArray(yq, dims=("yq",)),
+        "geolon":xr.DataArray(lon, dims=("yh", "xh")),
+        "geolat":xr.DataArray(lat, dims=("yh", "xh")),
+        "geolon_c":xr.DataArray(lon_c, dims=("yq", "xq",)),
+        "geolat_c":xr.DataArray(lat_c, dims=("yq", "xq",))
+    })
+    coords = {
+        'X': {'outer': 'xq', 'center': 'xh'},
+        'Y': {'outer': 'yq', 'center': 'yh'}
+    }
+    boundary = {
+        'X': "extend", # Specifying None causes this to default to `periodic` for some reason
+        'Y': "extend"  # Specifying None causes this to default to `periodic` for some reason
+    }
+    grid = xgcm.Grid(ds, coords=coords, boundary=boundary, autoparse_metadata=False)
+    return grid
+
+def test_convergent_transport():
+    from sectionate.section import grid_section
+    from sectionate.transports import convergent_transport
+    grid = initialize_minimal_outer_grid()
+    grid._ds['u'] = xr.DataArray(np.array([[1., -np.sqrt(2.)]]), dims=("yh","xq",))
+    grid._ds['v'] = xr.DataArray(np.array([[0], [np.pi]]), dims=("yq","xh",))
+    
+    # closed path around the whole square domain
+    lonseg = np.array([0, 1, 1, 0, 0])
+    latseg = np.array([0, 0, 1, 1, 0])
+    i, j, lons, lats = grid_section(grid, lonseg, latseg)
+
+    conv = convergent_transport(
+        grid,
+        i,
+        j,
+        utr="u",
+        vtr="v",
+        layer=None,
+        geometry="cartesian"
+    )['conv_mass_transport'].sum().values
+    
+    assert np.isclose(1. + 0. + np.sqrt(2.) - np.pi, conv, rtol=1.e-14)
+    
+    
+def initialize_minimal_spherical_grid():
+    xq = np.array([0., 120., 240., 360.])
+    yq = np.array([-80, 0., 80.])
+    xh = np.array([60., 180., 300.])
+    yh = np.array([-40., 40.])
+
+    lon, lat = np.meshgrid(xh, yh)
+    lon_c, lat_c = np.meshgrid(xq, yq)
+    ds = xr.Dataset({}, coords={
+        "xh":xr.DataArray(xh, dims=("xh",)),
+        "yh":xr.DataArray(yh, dims=("yh",)),
+        "xq":xr.DataArray(xq, dims=("xq",)),
+        "yq":xr.DataArray(yq, dims=("yq",)),
+        "geolon":xr.DataArray(lon, dims=("yh", "xh")),
+        "geolat":xr.DataArray(lat, dims=("yh", "xh")),
+        "geolon_c":xr.DataArray(lon_c, dims=("yq", "xq",)),
+        "geolat_c":xr.DataArray(lat_c, dims=("yq", "xq",))
+    })
+    coords = {
+        'X': {'outer': 'xq', 'center': 'xh'},
+        'Y': {'outer': 'yq', 'center': 'yh'}
+    }
+    boundary = {
+        'X': "periodic",
+        'Y': "extend"
+    }
+    grid = xgcm.Grid(ds, coords=coords, boundary=boundary, autoparse_metadata=False)
+    return grid
+    
+def test_convergent_transport_convention():
+    from sectionate.section import grid_section
+    from sectionate.transports import convergent_transport
+    grid = initialize_minimal_spherical_grid()
+    u = np.zeros((grid._ds.yh.size, grid._ds.xq.size))
+    v = np.ones((grid._ds.yq.size, grid._ds.xh.size))
+    grid._ds['u'] = xr.DataArray(u, dims=("yh", "xq"))
+    grid._ds['v'] = xr.DataArray(v, dims=("yq", "xh"))
+
+    lonseg = np.array([0., -120., -240., -360.])
+    latseg = np.array([0., 0., 0., 0.])
+    
+    i, j, lons, lats = grid_section(grid, lonseg, latseg)
+    conv = convergent_transport(
+        grid,
+        i,
+        j,
+        utr="u",
+        vtr="v",
+        layer=None
+    )['conv_mass_transport'].sum().values
+    
+    conv_rev = convergent_transport(
+        grid,
+        i[::-1],
+        j[::-1],
+        utr="u",
+        vtr="v",
+        layer=None
+    )['conv_mass_transport'].sum().values
+    
+    assert np.equal(-3., conv) and np.equal(-3, conv_rev)

sectionate/tests/test_section.py

@@ -0,0 +1,92 @@
+import numpy as np
+import xarray as xr
+
+
+# define simple lat-lon grid
+lon, lat = np.meshgrid(np.arange(360), np.arange(-80, 81))
+ds = xr.Dataset()
+ds["lon"] = xr.DataArray(lon, dims=("y", "x"))
+ds["lat"] = xr.DataArray(lat, dims=("y", "x"))
+
+def test_distance_on_unit_sphere():
+    from sectionate.section import distance_on_unit_sphere
+
+    # test of few points with unit radius
+    d = distance_on_unit_sphere(0, 0, 1.e-20, 0, R=1.)
+    assert np.isclose(d, 0., atol=1.e-14)
+    d = distance_on_unit_sphere(0, 0, 360, 0, R=1.)
+    assert np.isclose(d, 0., atol=1.e-14)
+    d = distance_on_unit_sphere(0, 45, 0, -45, R=1.)
+    assert np.isclose(d, np.pi/2, atol=1.e-14)
+    d = distance_on_unit_sphere(0, 0, 180, 0, R=1.)
+    assert np.isclose(d, np.pi, atol=1.e-14)
+    d = distance_on_unit_sphere(180, 0, 90, 0, R=1.)
+    assert np.isclose(d, np.pi/2, atol=1.e-14)
+    d = distance_on_unit_sphere(180, 45, 180, 0, R=1.)
+    assert np.isclose(d, np.pi/4, atol=1.e-14)
+
+
+def test_find_closest_grid_point():
+    from sectionate.section import find_closest_grid_point
+
+    # check it works with numpy arrays
+    i, j = find_closest_grid_point(0, 0, lon, lat)
+    assert np.equal(i, 0)
+    assert np.equal(j, 80)
+
+    # and xarray
+    i, j = find_closest_grid_point(0, 0, ds["lon"], ds["lat"])
+    assert np.equal(i, 0)
+    assert np.equal(j, 80)
+
+    i, j = find_closest_grid_point(180, 80, ds["lon"], ds["lat"])
+    assert np.equal(i, 180)
+    assert np.equal(j, 160)
+
+
+def test_grid_path():
+    from sectionate.section import infer_grid_path
+
+    # test zonal line
+    isec, jsec, lonsec, latsec = infer_grid_path(0, 80, 179, 80, lon, lat)
+    assert len(isec) == 180
+    assert lonsec[0] == 0.0
+    assert lonsec[-1] == 179.0
+    assert latsec[0] == 0.0
+    assert latsec[-1] == 0.0
+
+    # test merid line
+    isec, jsec, lonsec, latsec = infer_grid_path(0, 0, 0, 160, lon, lat)
+    assert len(isec) == 161
+    assert lonsec[0] == 0.
+    assert lonsec[-1] == 0.
+    assert latsec[0] == -80.0
+    assert latsec[-1] == 80.0
+
+    # test diagonal
+    isec, jsec, lonsec, latsec = infer_grid_path(0, 0, 100, 100, lon, lat)
+    assert len(isec) == 201  # expect ni+nj+1 values
+    isec, jsec, lonsec, latsec = infer_grid_path(0, 0, 50, 100, lon, lat)
+    assert len(isec) == 151  # expect ni+nj+1 values
+    isec, jsec, lonsec, latsec = infer_grid_path(10, 10, 100, 50, lon, lat)
+    assert len(isec) == 131  # expect ni+nj+1 values
+
+
+def test_infer_grid_path_from_geo():
+    from sectionate.section import infer_grid_path_from_geo
+
+    # test zonal line
+    isec, jsec, lonsec, latsec = infer_grid_path_from_geo(0, 0, 179, 0, lon, lat)
+    assert len(isec) == 180
+    assert lonsec[0] == 0.0
+    assert lonsec[-1] == 179.0
+    assert latsec[0] == 0.0
+    assert latsec[-1] == 0.0
+
+    # test merid line
+    isec, jsec, lonsec, latsec = infer_grid_path_from_geo(180, -80, 180, 80, lon, lat)
+    assert len(isec) == 161
+    assert lonsec[0]  == 180.0
+    assert lonsec[-1] == 180.0
+    assert latsec[0]  == -80.0
+    assert latsec[-1] == 80.0

sectionate/tests/test_section_class.py

@@ -0,0 +1,59 @@
+import numpy as np
+import xarray as xr
+import xgcm
+
+# define simple xgcm grid
+xq = np.array([0., 60, 120, 180, 240, 300., 360.])
+yq = np.array([-80., -40, 0, 40, 80.])
+
+lon_c, lat_c = np.meshgrid(xq, yq)
+ds = xr.Dataset({}, coords={
+    "xq":xr.DataArray(xq, dims=("xq",)),
+    "yq":xr.DataArray(yq, dims=("yq",)),
+    "lon_c":xr.DataArray(lon_c, dims=("yq", "xq",)),
+    "lat_c":xr.DataArray(lat_c, dims=("yq", "xq",))
+})
+coords = {
+    'X': {'outer': 'xq'},
+    'Y': {'outer': 'yq'}
+}
+boundary = {
+    'X': 'periodic',
+    'Y': 'extend'
+}
+grid = xgcm.Grid(ds, coords=coords, boundary=boundary, autoparse_metadata=False)
+
+def modequal(a,b):
+    return np.equal(np.mod(a, 360.), np.mod(b, 360.))
+
+def test_open_gridded_section():
+    from sectionate.section import Section, GriddedSection
+    lonseg = np.array([0., 120, 120, 0])
+    latseg = np.array([-80., -80, 0, 0])
+    sec = Section("testsec", (lonseg, latseg))
+    sec_gridded = GriddedSection(sec, grid)
+
+    assert np.all([
+        modequal(sec_gridded.i_c, np.array([0, 1, 2, 2, 2, 1, 0])),
+        modequal(sec_gridded.j_c, np.array([0, 0, 0, 1, 2, 2, 2])),
+        modequal(sec_gridded.lons_c, np.array([0.,  60., 120., 120., 120.,  60., 0.])),
+        modequal(sec_gridded.lats_c, np.array([-80., -80., -80., -40.,   0.,   0.,   0.]))
+    ])
+
+def test_closed_gridded_parent_section():
+    from sectionate.section import Section, join_sections, GriddedSection
+    lonseg = np.array([  0., 120, 120,  0,   0])
+    latseg = np.array([-80., -80,   0,  0, -80.])
+    # Test join_sections and children/parent relationships
+    sec1 = Section("sec1", (lonseg[0:3], latseg[0:3]))
+    sec2 = Section("sec2", (lonseg[2: ], latseg[2: ]))
+    sec = join_sections("sec", sec1, sec2)
+    assert isinstance(sec.children["sec1"], Section)
+    # Test results from join_section
+    sec_gridded = GriddedSection(sec, grid)
+    assert np.all([
+        modequal(sec_gridded.i_c, np.array([0, 1, 2, 2, 2, 1, 0, 0, 0])),
+        modequal(sec_gridded.j_c, np.array([0, 0, 0, 1, 2, 2, 2, 1, 0])),
+        modequal(sec_gridded.lons_c, np.array([0.,  60., 120., 120., 120.,  60., 0., 0., 0.])),
+        modequal(sec_gridded.lats_c, np.array([-80., -80., -80., -40.,   0.,   0.,   0., -40., -80.]))
+    ])

sectionate/tests/test_section_cornercases.py

@@ -0,0 +1,63 @@
+import numpy as np
+import xarray as xr
+import xgcm
+
+# define simple xgcm grid
+xq = np.array([0., 60, 120, 180, 240, 300., 360.])
+yq = np.array([-80., -40, 0, 40, 80.])
+
+lon_c, lat_c = np.meshgrid(xq, yq)
+ds = xr.Dataset({}, coords={
+    "xq":xr.DataArray(xq, dims=("xq",)),
+    "yq":xr.DataArray(yq, dims=("yq",)),
+    "lon_c":xr.DataArray(lon_c, dims=("yq", "xq",)),
+    "lat_c":xr.DataArray(lat_c, dims=("yq", "xq",))
+})
+coords = {
+    'X': {'outer': 'xq'},
+    'Y': {'outer': 'yq'}
+}
+boundary = {
+    'X': 'periodic',
+    'Y': 'extend'
+}
+grid = xgcm.Grid(ds, coords=coords, boundary=boundary, autoparse_metadata=False)
+
+def modequal(a,b):
+    return np.equal(np.mod(a, 360.), np.mod(b, 360.))
+
+def test_open_grid_section():
+    from sectionate.section import grid_section
+    lonseg = np.array([0., 120, 120, 0])
+    latseg = np.array([-80., -80, 0, 0])
+    i, j, lons, lats = grid_section(grid, lonseg, latseg)
+    assert np.all([
+        modequal(i, np.array([0, 1, 2, 2, 2, 1, 0])),
+        modequal(j, np.array([0, 0, 0, 1, 2, 2, 2])),
+        modequal(lons, np.array([0.,  60., 120., 120., 120.,  60., 0.])),
+        modequal(lats, np.array([-80., -80., -80., -40.,   0.,   0.,   0.]))
+    ])
+    
+def test_closed_grid_section():
+    from sectionate.section import grid_section
+    lonseg = np.array([0., 120, 120, 0, 0])
+    latseg = np.array([-80., -80, 0, 0, -80.])
+    i, j, lons, lats = grid_section(grid, lonseg, latseg)
+    assert np.all([
+        modequal(i, np.array([0, 1, 2, 2, 2, 1, 0, 0, 0])),
+        modequal(j, np.array([0, 0, 0, 1, 2, 2, 2, 1, 0])),
+        modequal(lons, np.array([0.,  60., 120., 120., 120.,  60., 0., 0., 0.])),
+        modequal(lats, np.array([-80., -80., -80., -40.,   0.,   0.,   0., -40., -80.]))
+    ])
+    
+def test_periodic_grid_section():
+    from sectionate.section import grid_section
+    lonseg = np.array([300, 60])
+    latseg = np.array([0, 0])
+    i, j, lons, lats = grid_section(grid, lonseg, latseg)
+    assert np.all([
+        modequal(i, np.array([5, 0, 1])),
+        modequal(j, np.array([2, 2, 2])),
+        modequal(lons, np.array([300.,  0., 60.])),
+        modequal(lats, np.array([0.,   0.,   0.]))
+    ])

sectionate/tests/test_utils.py

@@ -0,0 +1,4 @@
+def test_load_section():
+    from sectionate.utils import get_all_section_names, load_section
+    section_names = get_all_section_names()
+    load_section(section_names[0])

sectionate/tracers.py

@@ -0,0 +1,76 @@
+import numpy as np
+import xarray as xr
+from .transports import (
+    uvindices_from_qindices
+)
+
+from .gridutils import (
+    check_symmetric,
+    coord_dict
+)
+
+def extract_tracer(
+    name,
+    grid,
+    i_c,
+    j_c,
+    sect_coord="sect"
+    ):
+    """
+    Extract tracer data on cell thickness grid along the grid path
+    of (i_c, j_c) for plotting.
+
+    PARAMETERS:
+    -----------
+    name:
+        name of variable in `grid._ds`
+    grid: xgcm.Grid
+        grid describing model and containing data
+    i_c: int
+        vorticity point indices along 'X' dimension 
+    j_c: int
+        vorticity point indices along 'Y' dimension
+    sect_coord: str
+        Name of the dimension describing along-section data in the output. Default: 'sect'.
+
+    RETURNS:
+    --------
+
+    xarray.DataArray with data interpolated to the U and V points along the section.
+    """
+    
+    da=grid._ds[name]
+    coords = coord_dict(grid)
+    symmetric = check_symmetric(grid)
+    
+    # get indices of UV points from broken line
+    uvindices = uvindices_from_qindices(grid, i_c, j_c)
+        
+    section = xr.Dataset()
+    section["i"] = xr.DataArray(uvindices["i"], dims=sect_coord)
+    section["j"] = xr.DataArray(uvindices["j"], dims=sect_coord)
+    section["Umask"] = xr.DataArray(uvindices["var"]=="U", dims=sect_coord)
+    section["Vmask"] = xr.DataArray(uvindices["var"]=="V", dims=sect_coord)
+
+    increment = 1 if symmetric else 0
+    usel_left  = {coords["X"]["center"]: np.mod(section["i"]-increment  , da[coords["X"]["center"]].size),
+                  coords["Y"]["center"]: np.mod(section["j"]            , da[coords["Y"]["center"]].size)}
+    usel_right = {coords["X"]["center"]: np.mod(section["i"]-increment+1, da[coords["X"]["center"]].size),
+                  coords["Y"]["center"]: np.mod(section["j"]            , da[coords["Y"]["center"]].size)}
+
+    vsel_left  = {coords["X"]["center"]: np.mod(section["i"]            , da[coords["X"]["center"]].size),
+                  coords["Y"]["center"]: np.mod(section["j"]-increment  , da[coords["Y"]["center"]].size)}
+    vsel_right = {coords["X"]["center"]: np.mod(section["i"]            , da[coords["X"]["center"]].size),
+                  coords["Y"]["center"]: np.mod(section["j"]-increment+1, da[coords["Y"]["center"]].size)}
+
+    tracer = sum([
+        xr.where(~np.isnan(da.isel(usel_right)), 0.5*da.isel(usel_left),  da.isel(usel_left) ).fillna(0.) * section["Umask"],
+        xr.where(~np.isnan(da.isel(usel_left )), 0.5*da.isel(usel_right), da.isel(usel_right)).fillna(0.) * section["Umask"],
+        xr.where(~np.isnan(da.isel(vsel_right)), 0.5*da.isel(vsel_left),  da.isel(vsel_left) ).fillna(0.) * section["Vmask"],
+        xr.where(~np.isnan(da.isel(vsel_left )), 0.5*da.isel(vsel_right), da.isel(vsel_right)).fillna(0.) * section["Vmask"],
+    ])
+    tracer = tracer.where(tracer!=0., np.nan)
+    tracer.name = da.name
+    tracer.attrs = da.attrs
+
+    return tracer