логическое назначение xarray python - PullRequest
Новая
       17

логическое назначение xarray python

0 голосов
/ 09 июня 2019

Мне нужно присвоить значения одного xr.DataArray (gt), если выполняется определенное условие. И значения другого xr.DataArray (lt), если выполняется другое условие. Я хочу, чтобы он масштабировался (может быть более двух условий, которые могут не охватывать все возможности).

В идеале я бы хотел сделать все это в xarray.

В настоящее время у меня есть решение, которое записывает в numpy и использует встроенную функциональность маскированного массива.

Вот минимальный воспроизводимый пример 

import pandas as pd
import numpy as np
import xarray as xr

time = pd.date_range('2010-01-01','2011-12-31',freq='M')
lat = np.linspace(-5.175003, -4.7250023, 10)
lon = np.linspace(33.524994, 33.97499, 10)
precip = np.random.normal(0, 1, size=(len(time), len(lat), len(lon)))

ds = xr.Dataset(
    {'precip': (['time', 'lat', 'lon'], precip)},
    coords={
        'lon': lon,
        'lat': lat,
        'time': time,
    }
)

# DUMMY DataArray's for filling in gt / lt values
lt = xr.full_like(da, -100)
gt = xr.full_like(da, 100)

что я хочу

Я хочу DataArray со значениями lt, когда da.values МЕНЬШЕ 1 ('-100'), и значениями gt (100), когда da.values БОЛЬШЕ, ЧЕМ 1. 

<xarray.Dataset>
Dimensions:  (lat: 10, lon: 10, time: 24)
Coordinates:
  * lon      (lon) float64 33.52 33.57 33.62 33.67 ... 33.82 33.87 33.92 33.97
  * lat      (lat) float64 -5.175 -5.125 -5.075 -5.025 ... -4.825 -4.775 -4.725
  * time     (time) datetime64[ns] 2010-01-31 2010-02-28 ... 2011-12-31
Data variables:
    out      (time, lat, lon) float64 -100.0 -100.0 -100.0 ... -100.0 -100.0

Что я пробовал:

  • с использованием синтаксиса .where для выбора другого (того же size xr.Dataset) объекта
  • преобразовать в np.ma.array
  • рекомбинировать, используя np.ma.array синтаксис
  • читать обратно в xr.Dataset
da = ds.precip

# get the overall mask (e.g. land-sea mask)
overall_mask = da.isnull()


# create masked `xr.DataArray` with values from `lt` or `gt`
lt1 = da.where(~(da < 1), other=lt)
lt_mask = lt1.where(~(da < 1)).isnull()
lt1 = lt1.where(lt_mask).where(~overall_mask)

gt1 = da.where(~(da >= 1), other=gt)
gt_mask = gt1.where(~(da >= 1)).isnull()
gt1 = gt1.where(gt_mask).where(~overall_mask)

# convert to numpy masked array
gtm_np = gt_mask.values
gt1_np = gt1.values
gt1_np = np.ma.array(gt1_np, mask=(~gtm_np))

ltm_np = lt_mask.values
lt1_np = lt1.values
lt1_np = np.ma.array(lt1_np, mask=(~ltm_np))

# recombine masked arrays
out_array = (
    np.ma.array(
        gt1_np.filled(1) * lt1_np.filled(1),
        mask=(gt1_np.mask * lt1_np.mask)
    )
)

# assign back to xarray
new = xr.ones_like(da).to_dataset()
new['out'] = (['time','lat','lon'], out_array)
new = new.drop('precip')
new

Это выглядит так:

Example output

Мне нужен способ сделать это эффективно / изначально в xarray, если это возможно.

...