Masking DataArrays and modifying values based on dimensions

[VeckoTheGecko]

Hey all,
I am wanting to modify values in a DataArray, but I'm struggling to find a way in xarray to create a mask using the dimensions, and then assign a value accordingly.

For example, I want my land_flag variable to be True for lon < 15 (I instantiated it to False to begin with).
I tried to do this, however it failed (note my_data is a dataset):

>>> my_data["land_flag"].where(my_data.lon < 15) = xr.full_like(my_data["land_flag"].where(my_data.lon < 15), True)
SyntaxError: cannot assign to function call here. Maybe you meant '==' instead of '='?

I know that I can sidestep this issue by feeding in a numpy array such that the flag is True for longitude coordinates < 15, but I feel like there should be a more intuitive/flexible approach that takes advantage of xarrays indexing.

Sorry if this is a simple question. I've been looking for a while but I haven't been able to find anything. Also, I'm new to xarray so any advice or tips are more than welcome 🙂 .

Full example code:

import xarray as xr
import numpy as np
import datetime

def iso_to_datetime(string):
    """Input a string in the form '2022-01-01' and it gets converted to a datetime object.
    """
    return datetime.datetime.fromisoformat(string)

def generate_dataset(
    lon_interval: tuple[float, float],
    lat_interval: tuple[float, float],
    time_interval: tuple[datetime.datetime, datetime.datetime], 
    degree_resolution: float,
    time_resolution: datetime.timedelta
    ) -> xr.Dataset:
    """
    Generates a synthetic xarray dataset for ocean surface currents given intervals for the dimensions and a resolution.
    """
    # Creating arrays for axes
    lon = np.arange(lon_interval[0], lon_interval[1], degree_resolution)
    lat = np.arange(lat_interval[0], lat_interval[1], degree_resolution)
    time = np.arange(time_interval[0], time_interval[1], time_resolution)

    # Calculating lengths of dimensions
    lon_length, lat_length, time_length = len(lon), len(lat), len(time)

    # Specifying the variables, their dimensions, data and attributes
    U_data = (
        ["lon", "lat", "time"], # Dimensions
        np.full((lon_length, lat_length, time_length), 0.5), # Data
        dict(units = "m/s", long_name = "current velocity east") # Setting attributes for U
        )
    
    V_data = (
        ["lon", "lat", "time"], # Dimensions
        np.full((lon_length, lat_length, time_length), 0), # Data
        dict(units = "m/s", long_name = "current velocity north") # Setting attributes for U
        )
    
    land_cells = int(0.5 * lon_length)
    land_flag_data  = (
        ["lon", "lat"], # Dimensions
        np.full((lon_length, lat_length), False),
        # This is the "fix" using numpy
        # np.r_[np.full((land_cells, lat_length), True), np.full((lon_length - land_cells, lat_length), False)], # Data, half land half ocean
        dict(units = "boolean", long_name = "boolean flag indicating whether a cell is land") # Setting attributes for U
    )

    dataset = xr.Dataset(
        data_vars = dict(
            U = U_data,
            V = V_data,
            land_flag = land_flag_data
        ),
        coords = dict(
            lon = lon,
            lat = lat,
            time = time,
        ),
        attrs = dict(
            long_name = "Idealised ocean current data",
            units = "m/s"
        )
    )
    dataset.lon.attrs = dict(units = "degrees_east", long_name = "longitude")
    dataset.lat.attrs = dict(units = "degrees_north", long_name = "latitude")
    dataset.time.attrs = dict(long_name = "time")

    return dataset

my_data = generate_dataset(
    lon_interval=(10,20),
    lat_interval=(20,30),
    time_interval=(iso_to_datetime("2022-01-01"), iso_to_datetime("2022-01-10")),
    degree_resolution=1/12,
    time_resolution=datetime.timedelta(hours = 3)
    )

my_data["land_flag"].where(my_data.lon < 15) = xr.full_like(my_data["land_flag"].where(my_data.lon < 15), True)

[dcherian]

Can you try the top level where function: https://docs.xarray.dev/en/stable/generated/xarray.where.html

land_flag = xr.where(ds.lon < 15, True, False)

[VeckoTheGecko]

Sorry for the late response, but thank you! This is exactly what I was looking for