How to create new columns and apply a function

[thomashirtz]

I have a dataset containing many different coord and data. I would like to take one of the dimension and apply a function to it.
For example:

import xarray as xr
import numpy as np

dataset = xr.Dataset(
    {"measurements": (("simulation", "characteristic", "value"), np.random.random((10,3,6)))},
    coords={"characteristic": ['a', 'b', 'c']})

I would like to select columns, for example b and c, and create a new column so that:

dataset['characteristic'] == ['a', 'b', 'c', 'log_b', 'log_c']
dataset['measurements'].values.shape == (10, 5, 6)

(I know that for example, with pandas, I can do something like df['log_b'] = np.log(df['b']))

I can easily "create" a new log dataset using :

columns = ['b', 'c']
np.log(xarray['measurements'].loc[:,columns])

but after I am not sure how to merge everything while keeping the rest of my dataset

Thanks!

[dcherian]

I would think of this operation as a concatentation. See https://xarray.pydata.org/en/stable/combining.html#concatenate

import numpy as np
import xarray as xr

dataset = xr.Dataset(
    {"measurements": (("simulation", "characteristic", "value"), np.random.random((10,3,6)))},
    coords={"characteristic": ['a', 'b', 'c']})

logged = np.log(dataset[["measurements"]].sel(characteristic=["b", "c"]))
logged = logged.assign_coords(characteristic=["log_b", "log_c"])

xr.concat([dataset, logged], dim="characteristic")

<xarray.Dataset>
Dimensions:         (characteristic: 5, simulation: 10, value: 6)
Coordinates:
  * characteristic  (characteristic) <U5 'a' 'b' 'c' 'log_b' 'log_c'
Dimensions without coordinates: simulation, value
Data variables:
    measurements    (simulation, characteristic, value) float64 0.9663 ... -0...

Do you have suggestions for improving our documentation on this point?

[thomashirtz]

I still have two issues, since my dataset possesses much more coords and data I get this error :

import numpy as np
import xarray as xr

dataset = xr.Dataset(
    {"measurements": (("simulation", "characteristic", "value"), np.random.random((10,3,6))),
    "more_data": (('a', 'b'), np.random.random((10,10)))},
    coords={"characteristic": ['a', 'b', 'c']})

logged = np.log(dataset[["measurements"]].sel(characteristic=["b", "c"]))
logged = logged.assign_coords(characteristic=["log_b", "log_c"])

xr.concat([dataset, logged], dim="characteristic")

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\dataset.py in _construct_dataarray(self, name)
   1304         try:
-> 1305             variable = self._variables[name]
   1306         except KeyError:

KeyError: 'more_data'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in _dataset_concat(datasets, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs)
    507             try:
--> 508                 vars = ensure_common_dims([ds[k].variable for ds in datasets])
    509             except KeyError:

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in <listcomp>(.0)
    507             try:
--> 508                 vars = ensure_common_dims([ds[k].variable for ds in datasets])
    509             except KeyError:

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\dataset.py in __getitem__(self, key)
   1408         if hashable(key):
-> 1409             return self._construct_dataarray(key)
   1410         else:

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\dataset.py in _construct_dataarray(self, name)
   1306         except KeyError:
-> 1307             _, name, variable = _get_virtual_variable(
   1308                 self._variables, name, self._level_coords, self.dims

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\dataset.py in _get_virtual_variable(variables, key, level_vars, dim_sizes)
    170     else:
--> 171         ref_var = variables[ref_name]
    172 

KeyError: 'more_data'

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-114-ddcab7d6c37c> in <module>
     10 logged = logged.assign_coords(characteristic=["log_b", "log_c"])
     11 
---> 12 xr.concat([dataset, logged], dim="characteristic")

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in concat(objs, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs)
    237             "objects, got %s" % type(first_obj)
    238         )
--> 239     return f(
    240         objs, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs
    241     )

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in _dataset_concat(datasets, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs)
    508                 vars = ensure_common_dims([ds[k].variable for ds in datasets])
    509             except KeyError:
--> 510                 raise ValueError("%r is not present in all datasets." % k)
    511             combined = concat_vars(vars, dim, positions)
    512             assert isinstance(combined, Variable)

ValueError: 'more_data' is not present in all datasets.

I also have this error and I am currently trying to generate a code that I can give that can reproduce it, but I didn't manage yet:

def insert_after_substring(string, string_to_insert, substring):
    index = string.find(substring)
    return string[:index+len(substring)] + string_to_insert + string[len(substring)+index:]
log_columns = [insert_chars_after_substring(c, '#', 'log_') for c in columns]

logged = np.log(dataset["measurements"].sel(characteristic=columns))
logged = logged.assign_coords(characteristic=log_columns)

xr.concat([dataset, logged], dim="characteristic")

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-125-7cabbec96a48> in <module>
      4 logged = logged.assign_coords(characteristic=log_columns)
      5 
----> 6 xr.concat([dataset, logged], dim="characteristic")

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in concat(objs, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs)
    237             "objects, got %s" % type(first_obj)
    238         )
--> 239     return f(
    240         objs, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs
    241     )

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in _dataset_concat(datasets, dim, data_vars, coords, compat, positions, fill_value, join, combine_attrs)
    433     )
    434 
--> 435     dim_coords, dims_sizes, coord_names, data_names = _parse_datasets(datasets)
    436     dim_names = set(dim_coords)
    437     unlabeled_dims = dim_names - coord_names

C:\ProgramData\Anaconda3\lib\site-packages\xarray\core\concat.py in _parse_datasets(datasets)
    394 
    395     for ds in datasets:
--> 396         dims_sizes.update(ds.dims)
    397         all_coord_names.update(ds.coords)
    398         data_vars.update(ds.data_vars)

ValueError: dictionary update sequence element #0 has length 10; 2 is required

(I checked for this error and the logged dataset and the dataset dataset seems okay, the simulation and the measurement_index match, I am confused by this error)

For the documentation, I tried to search into different docs, but it is quite hard to read examples when you didn't build it by yourself.. Maybe a repository of more examples that refer to data frames that you can also access by keywords ?
The problem is that this library have so many applications and the needs are so specific that it is hard to do such thing

[dcherian]

This terrible error is because dataset is a Dataset and logged is a DataArray. Use dataset[["measurements"]]. Can you open an issue for raising a better error in this case please?

[dcherian]

Ah for your first error you need

xr.concat([dataset, logged], dim="characteristic", data_vars="minimal")

"minimal" tells xarray to only concat variables with the "characteristic" dimension.

[thomashirtz]

Thank you very much! It is working well now :)

This terrible error is because dataset is a Dataset and logged is a DataArray. Use dataset[["measurements"]]. Can you open an issue for raising a better error in this case please?

I opened an issue here