Ops with scalars and vectors

CHANGELOG.rst

@@ -4,6 +4,7 @@ Changelog
 master
 ------
 
+- (`#82 <https://github.com/openscm/scmdata/pull/82>`_) Add support for adding Pint scalars and vectors to :obj:`scmdata.Timeseries` and :obj:`scmdata.ScmRun` instances
 - (`#85 <https://github.com/openscm/scmdata/pull/85>`_) Allow required columns to be read as ``extras`` from netCDF files (closes `#83 <https://github.com/openscm/scmdata/issues/83>`_)
 - (`#84 <https://github.com/openscm/scmdata/pull/84>`_) Raise a DeprecationWarning if no default ``inplace`` argument is provided for :meth:`ScmRun.drop_meta`. inplace default behaviour scheduled to be changed to ``False`` in v0.7.0
 - (`#81 <https://github.com/openscm/scmdata/pull/81>`_) Add :attr:`scmdata.run.ScmRun.metadata` to track :obj:`ScmRun` instance-specific metadata (closes `#77 <https://github.com/openscm/scmdata/issues/77>`_)

notebooks/netcdf.ipynb

@@ -2978,7 +2978,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.6"
+   "version": "3.7.4"
   }
  },
  "nbformat": 4,

src/scmdata/run.py

@@ -15,6 +15,7 @@
 
 import numpy as np
 import pandas as pd
+import pint
 import xarray as xr
 from dateutil import parser
 from xarray.core.ops import inject_binary_ops
@@ -603,8 +604,20 @@ def _binary_op(
         def func(self, other):
             if isinstance(other, ScmRun):
                 return NotImplemented
-            if not isinstance(other, numbers.Number) and len(other) != len(self):
-                raise ValueError("Incorrect length")
+
+            is_number = isinstance(other, (numbers.Number, pint.Quantity))
+            if not is_number:
+                other_ndim = len(other.shape)
+                if other_ndim == 1:
+                    if other.shape[0] != self.shape[1]:
+                        raise ValueError(
+                            "only vectors with the same number of timesteps "
+                            "as self ({}) are supported".format(self.shape[1])
+                        )
+                else:
+                    raise ValueError(
+                        "operations with {}d data are not supported".format(other_ndim)
+                    )
 
             ret = self.copy(copy_ts=False)
             ret._ts = [

src/scmdata/timeseries.py

@@ -10,6 +10,8 @@
 from typing import Any, Callable, List, Union
 
 import numpy as np
+import openscm_units.unit_registry as ur
+import pint
 import xarray as xr
 from xarray.core.ops import inject_binary_ops
 
@@ -205,12 +207,22 @@ def _binary_op(
         def func(self, other):
             other_data = getattr(other, "_data", other)
 
-            ts = (
-                f(self._data, other_data)
-                if not reflexive
-                else f(other_data, self._data)
-            )
+            if isinstance(other, pint.Quantity):
+                try:
+                    self_data = self._data * ur(self.meta["unit"])
+                except KeyError:
+                    # let Pint assume dimensionless and raise an error as
+                    # necessary
+                    self_data = self._data
+            else:
+                self_data = self._data
+
+            ts = f(self_data, other_data) if not reflexive else f(other_data, self_data)
             ts.attrs = self._data.attrs
+            if isinstance(other, pint.Quantity):
+                ts.attrs["unit"] = str(ts.data.units)
+                ts.data = ts.values
+
             return TimeSeries(ts)
 
         return func