transform xarray to 3D and export as tif file with rioxarray

Question

I have a multidimensional xarray. I want to select a specific value for the dimension called number and then get rid of it and only keep step, latitude and longitude dimensions and export the xarray as a tif file for a variable.

This is my starting xarray:

<xarray.Dataset>
Dimensions:            (number: 2, step: 2, latitude: 45, longitude: 45)
Coordinates:
  * number             (number) int32 1 2
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 06:00:00 12:00:00
    heightAboveGround  float64 ...
  * latitude           (latitude) float64 28.8 28.4 28.0 27.6 ... 12.0 11.6 11.2
  * longitude          (longitude) float64 111.6 112.0 112.4 ... 128.8 129.2
    valid_time         (step) datetime64[ns] dask.array<chunksize=(2,), meta=np.ndarray>
Data variables:
    v10                (number, step, latitude, longitude) float32 dask.array<chunksize=(2, 2, 15, 15), meta=np.ndarray>
    u10                (number, step, latitude, longitude) float32 dask.array<chunksize=(2, 2, 15, 15), meta=np.ndarray>
Attributes:
    GRIB_edition:            2
    GRIB_centre:             ecmf
    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             European Centre for Medium-Range Weather Forecasts
    history:                 2023-07-11T10:38 GRIB to CDM+CF via cfgrib-0.9.1...

Here I have masked the xarray so that only one number is returned and I have tried to squeeze the dimension but it's still there.

ensemble = 1
mask = (xarr.number.values == ensemble)
xarr1 = xarr.sel(number=mask).squeeze(drop=True)
xarr1.rio.write_crs("epsg:4326", inplace=True) # set wgs84 proj to xarray

<xarray.Dataset>
Dimensions:            (number: 0, step: 2, latitude: 45, longitude: 45)
Coordinates:
  * number             (number) int32 
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 06:00:00 12:00:00
    heightAboveGround  float64 ...
  * latitude           (latitude) float64 28.8 28.4 28.0 27.6 ... 12.0 11.6 11.2
  * longitude          (longitude) float64 111.6 112.0 112.4 ... 128.8 129.2
    valid_time         (step) datetime64[ns] dask.array<chunksize=(2,), meta=np.ndarray>
    spatial_ref        int32 0
Data variables:
    v10                (number, step, latitude, longitude) float32 dask.array<chunksize=(0, 2, 15, 15), meta=np.ndarray>
    u10                (number, step, latitude, longitude) float32 dask.array<chunksize=(0, 2, 15, 15), meta=np.ndarray>
Attributes:
    GRIB_edition:            2
    GRIB_centre:             ecmf
    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             European Centre for Medium-Range Weather Forecasts
    history:                 2023-07-11T10:56 GRIB to CDM+CF via cfgrib-0.9.1...

I calculate the wind speed and add it as a new variable called ws to the xarray

ds2 = xarr1.assign(ws=(xarr1["u10"] ** 2 + xarr1["v10"] ** 2) ** 0.5)

I create a new xarray containing only the wind speed variable

xarr_ws = ds2.ws

<xarray.DataArray 'ws' (number: 0, step: 2, latitude: 45, longitude: 45)>
dask.array<pow, shape=(0, 2, 45, 45), dtype=float32, chunksize=(0, 2, 24, 24), chunktype=numpy.ndarray>
Coordinates:
  * number             (number) int32 
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 06:00:00 12:00:00
    heightAboveGround  float64 ...
  * latitude           (latitude) float64 28.8 28.4 28.0 27.6 ... 12.0 11.6 11.2
  * longitude          (longitude) float64 111.6 112.0 112.4 ... 128.8 129.2
    valid_time         (step) datetime64[ns] dask.array<chunksize=(2,), meta=np.ndarray>
    spatial_ref        int32 0

Use rioxarray to export to a tif file but this fails because I have 4 dimensions rather than 3.

# save tif with wind speed with all timesteps
xarr_ws.rio.to_raster(f'{storm_name[1:]}_ens{ensemble}_wind_speed_by_timestep.tif')

error message:

in _check_dimensions raise TooManyDimensions(rioxarray.exceptions.TooManyDimensions: Only 2D and 3D data arrays supported. Data variable: ws

How do I reduce the dimensions from (number: 0, step: 2, latitude: 45, longitude: 45) to (step: 2, latitude: 45, longitude: 45) ?

I have tried .squeeze('number'), .drop(dim='number') and .drop_dims(number) with no luck